Add support for pytorch 2.6+

2d_classification/mednist_tutorial.ipynb

@@ -575,7 +575,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "y_true = []\n",
     "y_pred = []\n",

2d_segmentation/torch/unet_evaluation_array.py

@@ -58,7 +58,7 @@ def main(tempdir):
         num_res_units=2,
     ).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_segmentation2d_array.pth"))
+    model.load_state_dict(torch.load("best_metric_model_segmentation2d_array.pth", weights_only=True))
     model.eval()
     with torch.no_grad():
         for val_data in val_loader:

2d_segmentation/torch/unet_evaluation_dict.py

@@ -72,7 +72,7 @@ def main(tempdir):
         num_res_units=2,
     ).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_segmentation2d_dict.pth"))
+    model.load_state_dict(torch.load("best_metric_model_segmentation2d_dict.pth", weights_only=True))
 
     model.eval()
     with torch.no_grad():

3d_classification/torch/densenet_evaluation_array.py

@@ -57,7 +57,7 @@ def main():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = monai.networks.nets.DenseNet121(spatial_dims=3, in_channels=1, out_channels=2).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_classification3d_array.pth"))
+    model.load_state_dict(torch.load("best_metric_model_classification3d_array.pth", weights_only=True))
     model.eval()
     with torch.no_grad():
         num_correct = 0.0

3d_classification/torch/densenet_evaluation_dict.py

@@ -63,7 +63,7 @@ def main():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = monai.networks.nets.DenseNet121(spatial_dims=3, in_channels=1, out_channels=2).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_classification3d_dict.pth"))
+    model.load_state_dict(torch.load("best_metric_model_classification3d_dict.pth", weights_only=True))
     model.eval()
     with torch.no_grad():
         num_correct = 0.0

3d_registration/learn2reg_nlst_paired_lung_ct.ipynb

@@ -872,7 +872,7 @@
    "source": [
     "# Automatic mixed precision (AMP) for faster training\n",
     "amp_enabled = True\n",
-    "scaler = torch.cuda.amp.GradScaler()\n",
+    "scaler = torch.GradScaler(\"cuda\")\n",
     "\n",
     "# Tensorboard\n",
     "if do_save:\n",
@@ -1127,7 +1127,7 @@
     "    )\n",
     "    # load model weights\n",
     "    filename_best_model = glob.glob(os.path.join(dir_load, \"segresnet_kpt_loss_best_tre*\"))[0]\n",
-    "    model.load_state_dict(torch.load(filename_best_model))\n",
+    "    model.load_state_dict(torch.load(filename_best_model, weights_only=True))\n",
     "    # to GPU\n",
     "    model.to(device)\n",
     "\n",
@@ -1139,7 +1139,7 @@
     "# Forward pass\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
-    "    with torch.cuda.amp.autocast(enabled=amp_enabled):\n",
+    "    with torch.autocast(\"cuda\", enabled=amp_enabled):\n",
     "        ddf_image, ddf_keypoints, pred_image, pred_label = forward(\n",
     "            check_data[\"fixed_image\"].to(device),\n",
     "            check_data[\"moving_image\"].to(device),\n",

3d_registration/learn2reg_oasis_unpaired_brain_mr.ipynb

@@ -610,7 +610,7 @@
    "source": [
     "# Automatic mixed precision (AMP) for faster training\n",
     "amp_enabled = True\n",
-    "scaler = torch.cuda.amp.GradScaler()\n",
+    "scaler = torch.GradScaler(\"cuda\")\n",
     "\n",
     "# Tensorboard\n",
     "if do_save:\n",
@@ -646,7 +646,7 @@
     "\n",
     "        # Forward pass and loss\n",
     "        optimizer.zero_grad()\n",
-    "        with torch.cuda.amp.autocast(enabled=amp_enabled):\n",
+    "        with torch.autocast(\"cuda\", enabled=amp_enabled):\n",
     "            ddf_image, pred_image, pred_label_one_hot = forward(\n",
     "                fixed_image, moving_image, moving_label, model, warp_layer, num_classes=4\n",
     "            )\n",
@@ -694,7 +694,7 @@
     "                # moving_label_35 = batch_data[\"moving_label_35\"].to(device)\n",
     "                n_steps += 1\n",
     "                # Infer\n",
-    "                with torch.cuda.amp.autocast(enabled=amp_enabled):\n",
+    "                with torch.autocast(\"cuda\", enabled=amp_enabled):\n",
     "                    ddf_image, pred_image, pred_label_one_hot = forward(\n",
     "                        fixed_image, moving_image, moving_label_4, model, warp_layer, num_classes=4\n",
     "                    )\n",
@@ -840,7 +840,7 @@
     "    model = VoxelMorph()\n",
     "    # load model weights\n",
     "    filename_best_model = glob.glob(os.path.join(dir_load, \"voxelmorph_loss_best_dice_*\"))[0]\n",
-    "    model.load_state_dict(torch.load(filename_best_model))\n",
+    "    model.load_state_dict(torch.load(filename_best_model, weights_only=True))\n",
     "    # to GPU\n",
     "    model.to(device)\n",
     "\n",
@@ -860,7 +860,7 @@
     "# Forward pass\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
-    "    with torch.cuda.amp.autocast(enabled=amp_enabled):\n",
+    "    with torch.autocast(\"cuda\", enabled=amp_enabled):\n",
     "        ddf_image, pred_image, pred_label_one_hot = forward(\n",
     "            fixed_image, moving_image, moving_label_35, model, warp_layer, num_classes=35\n",
     "        )"

3d_registration/paired_lung_ct.ipynb

@@ -860,7 +860,7 @@
     "resource = \"https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/pair_lung_ct.pth\"\n",
     "dst = f\"{root_dir}/pretrained_weight.pth\"\n",
     "download_url(resource, dst)\n",
-    "model.load_state_dict(torch.load(dst))"
+    "model.load_state_dict(torch.load(dst, weights_only=True))"
    ]
   },
   {

3d_segmentation/brats_segmentation_3d.ipynb

@@ -473,14 +473,14 @@
     "        )\n",
     "\n",
     "    if VAL_AMP:\n",
-    "        with torch.cuda.amp.autocast():\n",
+    "        with torch.autocast(\"cuda\"):\n",
     "            return _compute(input)\n",
     "    else:\n",
     "        return _compute(input)\n",
     "\n",
     "\n",
     "# use amp to accelerate training\n",
-    "scaler = torch.cuda.amp.GradScaler()\n",
+    "scaler = torch.GradScaler(\"cuda\")\n",
     "# enable cuDNN benchmark\n",
     "torch.backends.cudnn.benchmark = True"
    ]
@@ -526,7 +526,7 @@
     "            batch_data[\"label\"].to(device),\n",
     "        )\n",
     "        optimizer.zero_grad()\n",
-    "        with torch.cuda.amp.autocast():\n",
+    "        with torch.autocast(\"cuda\"):\n",
     "            outputs = model(inputs)\n",
     "            loss = loss_function(outputs, labels)\n",
     "        scaler.scale(loss).backward()\n",
@@ -733,7 +733,7 @@
     }
    ],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
     "    # select one image to evaluate and visualize the model output\n",
@@ -835,7 +835,7 @@
     }
    ],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "\n",
     "with torch.no_grad():\n",
@@ -924,7 +924,7 @@
     "        )\n",
     "\n",
     "    if VAL_AMP:\n",
-    "        with torch.cuda.amp.autocast():\n",
+    "        with torch.autocast(\"cuda\"):\n",
     "            return _compute(input)\n",
     "    else:\n",
     "        return _compute(input)"
@@ -977,7 +977,7 @@
    "source": [
     "onnx_model_path = os.path.join(root_dir, \"best_metric_model.onnx\")\n",
     "ort_session = onnxruntime.InferenceSession(onnx_model_path)\n",
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "\n",
     "with torch.no_grad():\n",

3d_segmentation/challenge_baseline/run_net.py

@@ -219,7 +219,7 @@ def infer(data_folder=".", model_folder="runs", prediction_folder="output"):
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     net = get_net().to(device)
-    net.load_state_dict(torch.load(ckpt, map_location=device))
+    net.load_state_dict(torch.load(ckpt, map_location=device, weights_only=True))
     net.eval()
 
     image_folder = os.path.abspath(data_folder)

3d_segmentation/spleen_segmentation_3d.ipynb

@@ -640,7 +640,7 @@
     }
    ],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
     "    for i, val_data in enumerate(val_loader):\n",
@@ -730,7 +730,7 @@
     }
    ],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "\n",
     "with torch.no_grad():\n",
@@ -827,7 +827,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "\n",
     "with torch.no_grad():\n",

3d_segmentation/spleen_segmentation_3d_visualization_basic.ipynb

@@ -823,7 +823,7 @@
     }
    ],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "\n",
     "with torch.no_grad():\n",

3d_segmentation/swin_unetr_brats21_segmentation_3d.ipynb

@@ -885,7 +885,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"model.pt\"))[\"state_dict\"])\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"model.pt\"), weights_only=True)[\"state_dict\"])\n",
     "model.to(device)\n",
     "model.eval()\n",
     "\n",

3d_segmentation/swin_unetr_btcv_segmentation_3d.ipynb

@@ -472,7 +472,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "weight = torch.load(\"./model_swinvit.pt\")\n",
+    "weight = torch.load(\"./model_swinvit.pt\", weights_only=True)\n",
     "model.load_from(weights=weight)\n",
     "print(\"Using pretrained self-supervied Swin UNETR backbone weights !\")"
    ]
@@ -493,7 +493,7 @@
     "torch.backends.cudnn.benchmark = True\n",
     "loss_function = DiceCELoss(to_onehot_y=True, softmax=True)\n",
     "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)\n",
-    "scaler = torch.cuda.amp.GradScaler()"
+    "scaler = torch.GradScaler(\"cuda\")"
    ]
   },
   {
@@ -516,7 +516,7 @@
     "    with torch.no_grad():\n",
     "        for batch in epoch_iterator_val:\n",
     "            val_inputs, val_labels = (batch[\"image\"].cuda(), batch[\"label\"].cuda())\n",
-    "            with torch.cuda.amp.autocast():\n",
+    "            with torch.autocast(\"cuda\"):\n",
     "                val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 4, model)\n",
     "            val_labels_list = decollate_batch(val_labels)\n",
     "            val_labels_convert = [post_label(val_label_tensor) for val_label_tensor in val_labels_list]\n",
@@ -537,7 +537,7 @@
     "    for step, batch in enumerate(epoch_iterator):\n",
     "        step += 1\n",
     "        x, y = (batch[\"image\"].cuda(), batch[\"label\"].cuda())\n",
-    "        with torch.cuda.amp.autocast():\n",
+    "        with torch.autocast(\"cuda\"):\n",
     "            logit_map = model(x)\n",
     "            loss = loss_function(logit_map, y)\n",
     "        scaler.scale(loss).backward()\n",
@@ -590,7 +590,7 @@
     "metric_values = []\n",
     "while global_step < max_iterations:\n",
     "    global_step, dice_val_best, global_step_best = train(global_step, train_loader, dice_val_best, global_step_best)\n",
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))"
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))"
    ]
   },
   {
@@ -679,7 +679,7 @@
    ],
    "source": [
     "case_num = 4\n",
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
     "    img_name = os.path.split(val_ds[case_num][\"image\"].meta[\"filename_or_obj\"])[1]\n",

3d_segmentation/torch/unet_evaluation_array.py

@@ -63,7 +63,7 @@ def main(tempdir):
         num_res_units=2,
     ).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_segmentation3d_array.pth"))
+    model.load_state_dict(torch.load("best_metric_model_segmentation3d_array.pth", weights_only=True))
     model.eval()
     with torch.no_grad():
         for val_data in val_loader:

3d_segmentation/torch/unet_evaluation_dict.py

@@ -81,7 +81,7 @@ def main(tempdir):
         num_res_units=2,
     ).to(devices[0])
 
-    model.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth"))
+    model.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth", weights_only=True))
 
     # if we have multiple GPUs, set data parallel to execute sliding window inference
     if len(devices) > 1:

3d_segmentation/torch/unet_inference_dict.py

@@ -91,7 +91,7 @@ def main(tempdir):
         strides=(2, 2, 2, 2),
         num_res_units=2,
     ).to(device)
-    net.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth"))
+    net.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth", weights_only=True))
 
     net.eval()
     with torch.no_grad():

3d_segmentation/unetr_btcv_segmentation_3d.ipynb

@@ -680,7 +680,7 @@
     "metric_values = []\n",
     "while global_step < max_iterations:\n",
     "    global_step, dice_val_best, global_step_best = train(global_step, train_loader, dice_val_best, global_step_best)\n",
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))"
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))"
    ]
   },
   {
@@ -769,7 +769,7 @@
    ],
    "source": [
     "case_num = 4\n",
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
     "    img_name = os.path.split(val_ds[case_num][\"image\"].meta[\"filename_or_obj\"])[1]\n",

acceleration/TensorRT_inference_acceleration.ipynb

@@ -284,7 +284,7 @@
     "device = workflow.device\n",
     "spatial_shape = (1, 3, 736, 480)\n",
     "model = workflow.network_def\n",
-    "model.load_state_dict(torch.load(model_weight))\n",
+    "model.load_state_dict(torch.load(model_weight, weights_only=True))\n",
     "model.to(device)\n",
     "model.eval()\n",
     "\n",

acceleration/automatic_mixed_precision.ipynb

@@ -289,7 +289,7 @@
     "    ).to(device)\n",
     "    loss_function = DiceLoss(to_onehot_y=True, softmax=True)\n",
     "    optimizer = torch.optim.Adam(model.parameters(), 1e-4)\n",
-    "    scaler = torch.cuda.amp.GradScaler() if amp else None\n",
+    "    scaler = torch.GradScaler(\"cuda\") if amp else None\n",
     "\n",
     "    post_pred = Compose([AsDiscrete(argmax=True, to_onehot=2)])\n",
     "    post_label = Compose([AsDiscrete(to_onehot=2)])\n",
@@ -321,7 +321,7 @@
     "            )\n",
     "            optimizer.zero_grad()\n",
     "            if amp and scaler is not None:\n",
-    "                with torch.cuda.amp.autocast():\n",
+    "                with torch.autocast(\"cuda\"):\n",
     "                    outputs = model(inputs)\n",
     "                    loss = loss_function(outputs, labels)\n",
     "                scaler.scale(loss).backward()\n",
@@ -353,7 +353,7 @@
     "                    roi_size = (160, 160, 128)\n",
     "                    sw_batch_size = 4\n",
     "                    if amp:\n",
-    "                        with torch.cuda.amp.autocast():\n",
+    "                        with torch.autocast(\"cuda\"):\n",
     "                            val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)\n",
     "                    else:\n",
     "                        val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)\n",