def METHOD_NAME():
    import torchviz
    dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = Unet3D()
    model.to(dev)
    model.eval()
    print(next(model.parameters()).is_cuda)  # True
    img = torch.randn(1, SIZE, SIZE, SIZE, 1).to(dev)
    out = model(img)
    dot = torchviz.make_dot(out, params=dict(model.named_parameters()), show_attrs=True, show_saved=True)
    dot.render("unet", format="png")
    torch.save(model, "model.pth")
    print(dot)