monai-weekly 1.5.dev2509__py3-none-any.whl → 1.5.dev2511__py3-none-any.whl

tests/inferers/test_latent_diffusion_inferer.py

@@ -19,7 +19,7 @@ from parameterized import parameterized
 
 from monai.inferers import LatentDiffusionInferer
 from monai.networks.nets import VQVAE, AutoencoderKL, DiffusionModelUNet, SPADEAutoencoderKL, SPADEDiffusionModelUNet
-from monai.networks.schedulers import DDPMScheduler
+from monai.networks.schedulers import DDPMScheduler, RFlowScheduler
 from monai.utils import optional_import
 
 _, has_einops = optional_import("einops")
@@ -339,31 +339,32 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
 
         input = torch.randn(input_shape).to(device)
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
-        timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],), device=input.device).long()
 
-        if dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+            timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],), device=input.device).long()
+
+            if dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                prediction = inferer(
+                    inputs=input,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    seg=input_seg,
+                    noise=noise,
+                    timesteps=timesteps,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            prediction = inferer(
-                inputs=input,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                seg=input_seg,
-                noise=noise,
-                timesteps=timesteps,
-            )
-        else:
-            prediction = inferer(
-                inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
-            )
-        self.assertEqual(prediction.shape, latent_shape)
+                prediction = inferer(
+                    inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
+                )
+            self.assertEqual(prediction.shape, latent_shape)
 
     @parameterized.expand(TEST_CASES)
     @skipUnless(has_einops, "Requires einops")
@@ -388,29 +389,30 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         stage_2.eval()
 
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
 
-        if ae_model_type == "SPADEAutoencoderKL" or dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            if ae_model_type == "SPADEAutoencoderKL" or dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                sample = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    seg=input_seg,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            sample = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                seg=input_seg,
-            )
-        else:
-            sample = inferer.sample(
-                input_noise=noise, autoencoder_model=stage_1, diffusion_model=stage_2, scheduler=scheduler
-            )
-        self.assertEqual(sample.shape, input_shape)
+                sample = inferer.sample(
+                    input_noise=noise, autoencoder_model=stage_1, diffusion_model=stage_2, scheduler=scheduler
+                )
+            self.assertEqual(sample.shape, input_shape)
 
     @parameterized.expand(TEST_CASES)
     @skipUnless(has_einops, "Requires einops")
@@ -437,37 +439,38 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         stage_2.eval()
 
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
 
-        if ae_model_type == "SPADEAutoencoderKL" or dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            if ae_model_type == "SPADEAutoencoderKL" or dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                sample, intermediates = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    seg=input_seg,
+                    save_intermediates=True,
+                    intermediate_steps=1,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            sample, intermediates = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                seg=input_seg,
-                save_intermediates=True,
-                intermediate_steps=1,
-            )
-        else:
-            sample, intermediates = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                save_intermediates=True,
-                intermediate_steps=1,
-            )
-        self.assertEqual(len(intermediates), 10)
-        self.assertEqual(intermediates[0].shape, input_shape)
+                sample, intermediates = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    save_intermediates=True,
+                    intermediate_steps=1,
+                )
+            self.assertEqual(len(intermediates), 10)
+            self.assertEqual(intermediates[0].shape, input_shape)
 
     @parameterized.expand(TEST_CASES)
     @skipUnless(has_einops, "Requires einops")
@@ -614,40 +617,40 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         conditioning_shape[1] = n_concat_channel
         conditioning = torch.randn(conditioning_shape).to(device)
 
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
-
-        timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],), device=input.device).long()
-
-        if dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],), device=input.device).long()
+
+            if dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                prediction = inferer(
+                    inputs=input,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    noise=noise,
+                    timesteps=timesteps,
+                    condition=conditioning,
+                    mode="concat",
+                    seg=input_seg,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            prediction = inferer(
-                inputs=input,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                noise=noise,
-                timesteps=timesteps,
-                condition=conditioning,
-                mode="concat",
-                seg=input_seg,
-            )
-        else:
-            prediction = inferer(
-                inputs=input,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                noise=noise,
-                timesteps=timesteps,
-                condition=conditioning,
-                mode="concat",
-            )
-        self.assertEqual(prediction.shape, latent_shape)
+                prediction = inferer(
+                    inputs=input,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    noise=noise,
+                    timesteps=timesteps,
+                    condition=conditioning,
+                    mode="concat",
+                )
+            self.assertEqual(prediction.shape, latent_shape)
 
     @parameterized.expand(TEST_CASES)
     @skipUnless(has_einops, "Requires einops")
@@ -681,36 +684,36 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         conditioning_shape[1] = n_concat_channel
         conditioning = torch.randn(conditioning_shape).to(device)
 
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
-
-        if dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            if dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                sample = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    conditioning=conditioning,
+                    mode="concat",
+                    seg=input_seg,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            sample = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                conditioning=conditioning,
-                mode="concat",
-                seg=input_seg,
-            )
-        else:
-            sample = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                conditioning=conditioning,
-                mode="concat",
-            )
-        self.assertEqual(sample.shape, input_shape)
+                sample = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    conditioning=conditioning,
+                    mode="concat",
+                )
+            self.assertEqual(sample.shape, input_shape)
 
     @parameterized.expand(TEST_CASES_DIFF_SHAPES)
     @skipUnless(has_einops, "Requires einops")
@@ -738,39 +741,39 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
 
         input = torch.randn(input_shape).to(device)
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        # We infer the VAE shape
-        autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]) - 1)) for i in input_shape[2:]]
-        inferer = LatentDiffusionInferer(
-            scheduler=scheduler,
-            scale_factor=1.0,
-            ldm_latent_shape=list(latent_shape[2:]),
-            autoencoder_latent_shape=autoencoder_latent_shape,
-        )
-        scheduler.set_timesteps(num_inference_steps=10)
-
-        timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],), device=input.device).long()
-
-        if dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
-            else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            prediction = inferer(
-                inputs=input,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                noise=noise,
-                timesteps=timesteps,
-                seg=input_seg,
-            )
-        else:
-            prediction = inferer(
-                inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            # We infer the VAE shape
+            autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]) - 1)) for i in input_shape[2:]]
+            inferer = LatentDiffusionInferer(
+                scheduler=scheduler,
+                scale_factor=1.0,
+                ldm_latent_shape=list(latent_shape[2:]),
+                autoencoder_latent_shape=autoencoder_latent_shape,
             )
-        self.assertEqual(prediction.shape, latent_shape)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],), device=input.device).long()
+
+            if dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                prediction = inferer(
+                    inputs=input,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    noise=noise,
+                    timesteps=timesteps,
+                    seg=input_seg,
+                )
+            else:
+                prediction = inferer(
+                    inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
+                )
+            self.assertEqual(prediction.shape, latent_shape)
 
     @parameterized.expand(TEST_CASES_DIFF_SHAPES)
     @skipUnless(has_einops, "Requires einops")
@@ -797,40 +800,42 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         stage_2.eval()
 
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        # We infer the VAE shape
-        if ae_model_type == "VQVAE":
-            autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]))) for i in input_shape[2:]]
-        else:
-            autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]) - 1)) for i in input_shape[2:]]
-
-        inferer = LatentDiffusionInferer(
-            scheduler=scheduler,
-            scale_factor=1.0,
-            ldm_latent_shape=list(latent_shape[2:]),
-            autoencoder_latent_shape=autoencoder_latent_shape,
-        )
-        scheduler.set_timesteps(num_inference_steps=10)
-
-        if dm_model_type == "SPADEDiffusionModelUNet" or ae_model_type == "SPADEAutoencoderKL":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            # We infer the VAE shape
+            if ae_model_type == "VQVAE":
+                autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]))) for i in input_shape[2:]]
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            prediction, _ = inferer.sample(
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                input_noise=noise,
-                save_intermediates=True,
-                seg=input_seg,
-            )
-        else:
-            prediction = inferer.sample(
-                autoencoder_model=stage_1, diffusion_model=stage_2, input_noise=noise, save_intermediates=False
+                autoencoder_latent_shape = [
+                    i // (2 ** (len(autoencoder_params["channels"]) - 1)) for i in input_shape[2:]
+                ]
+
+            inferer = LatentDiffusionInferer(
+                scheduler=scheduler,
+                scale_factor=1.0,
+                ldm_latent_shape=list(latent_shape[2:]),
+                autoencoder_latent_shape=autoencoder_latent_shape,
             )
-        self.assertEqual(prediction.shape, input_shape)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            if dm_model_type == "SPADEDiffusionModelUNet" or ae_model_type == "SPADEAutoencoderKL":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                prediction, _ = inferer.sample(
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    input_noise=noise,
+                    save_intermediates=True,
+                    seg=input_seg,
+                )
+            else:
+                prediction = inferer.sample(
+                    autoencoder_model=stage_1, diffusion_model=stage_2, input_noise=noise, save_intermediates=False
+                )
+            self.assertEqual(prediction.shape, input_shape)
 
     @skipUnless(has_einops, "Requires einops")
     def test_incompatible_spade_setup(self):
@@ -866,18 +871,19 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         stage_2.eval()
         noise = torch.randn((1, 3, 4, 4)).to(device)
         input_seg = torch.randn((1, 3, 8, 8)).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
 
-        with self.assertRaises(ValueError):
-            _ = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                seg=input_seg,
-            )
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            with self.assertRaises(ValueError):
+                _ = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    seg=input_seg,
+                )
 
 
 if __name__ == "__main__":

tests/integration/test_integration_bundle_run.py

@@ -76,8 +76,7 @@ class TestBundleRun(unittest.TestCase):
             )
         with open(meta_file, "w") as f:
             json.dump(
-                {"version": "0.1.0", "monai_version": "1.1.0", "pytorch_version": "1.13.1", "numpy_version": "1.22.2"},
-                f,
+                {"version": "0.1.0", "monai_version": "1.1.0", "pytorch_version": "2.3.0", "numpy_version": "1.22.2"}, f
             )
         cmd = ["coverage", "run", "-m", "monai.bundle"]
         # test both CLI entry "run" and "run_workflow"
@@ -114,8 +113,7 @@ class TestBundleRun(unittest.TestCase):
             )
         with open(meta_file, "w") as f:
             json.dump(
-                {"version": "0.1.0", "monai_version": "1.1.0", "pytorch_version": "1.13.1", "numpy_version": "1.22.2"},
-                f,
+                {"version": "0.1.0", "monai_version": "1.1.0", "pytorch_version": "2.3.0", "numpy_version": "1.22.2"}, f
             )
 
         os.mkdir(scripts_dir)

tests/integration/test_integration_classification_2d.py

@@ -166,7 +166,7 @@ def run_inference_test(root_dir, test_x, test_y, device="cuda:0", num_workers=10
     model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=len(np.unique(test_y))).to(device)
 
     model_filename = os.path.join(root_dir, "best_metric_model.pth")
-    model.load_state_dict(torch.load(model_filename))
+    model.load_state_dict(torch.load(model_filename, weights_only=True))
     y_true = []
     y_pred = []
     with eval_mode(model):

tests/integration/test_integration_fast_train.py

@@ -186,7 +186,7 @@ class IntegrationFastTrain(DistTestCase):
                 step += 1
                 optimizer.zero_grad()
                 # set AMP for training
-                with torch.cuda.amp.autocast():
+                with torch.autocast("cuda"):
                     outputs = model(batch_data["image"])
                     loss = loss_function(outputs, batch_data["label"])
                 scaler.scale(loss).backward()
@@ -207,7 +207,7 @@ class IntegrationFastTrain(DistTestCase):
                         roi_size = (96, 96, 96)
                         sw_batch_size = 4
                         # set AMP for validation
-                        with torch.cuda.amp.autocast():
+                        with torch.autocast("cuda"):
                             val_outputs = sliding_window_inference(val_data["image"], roi_size, sw_batch_size, model)
 
                         val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]

tests/integration/test_integration_segmentation_3d.py

@@ -216,7 +216,7 @@ def run_inference_test(root_dir, device="cuda:0"):
     ).to(device)
 
     model_filename = os.path.join(root_dir, "best_metric_model.pth")
-    model.load_state_dict(torch.load(model_filename))
+    model.load_state_dict(torch.load(model_filename, weights_only=True))
     with eval_mode(model):
         # resampling with align_corners=True or dtype=float64 will generate
         # slight different results between PyTorch 1.5 an 1.6

tests/metrics/test_compute_multiscalessim_metric.py

@@ -32,7 +32,7 @@ class TestMultiScaleSSIMMetric(unittest.TestCase):
         metric(preds, target)
         result = metric.aggregate()
         expected_value = 0.023176
-        self.assertTrue(expected_value - result.item() < 0.000001)
+        self.assertAlmostEqual(expected_value, result.item(), 4)
 
     def test2d_uniform(self):
         set_determinism(0)
@@ -45,7 +45,7 @@ class TestMultiScaleSSIMMetric(unittest.TestCase):
         metric(preds, target)
         result = metric.aggregate()
         expected_value = 0.022655
-        self.assertTrue(expected_value - result.item() < 0.000001)
+        self.assertAlmostEqual(expected_value, result.item(), 4)
 
     def test3d_gaussian(self):
         set_determinism(0)
@@ -58,7 +58,7 @@ class TestMultiScaleSSIMMetric(unittest.TestCase):
         metric(preds, target)
         result = metric.aggregate()
         expected_value = 0.061796
-        self.assertTrue(expected_value - result.item() < 0.000001)
+        self.assertAlmostEqual(expected_value, result.item(), 4)
 
     def input_ill_input_shape2d(self):
         metric = MultiScaleSSIMMetric(spatial_dims=3, weights=[0.5, 0.5])

tests/metrics/test_surface_dice.py

@@ -82,7 +82,7 @@ class TestAllSurfaceDiceMetrics(unittest.TestCase):
         expected_res0[1, 1] = np.nan
         for b, c in np.ndindex(batch_size, n_class):
             np.testing.assert_allclose(expected_res0[b, c], res0[b, c].cpu())
-        np.testing.assert_array_equal(agg0.cpu(), np.nanmean(np.nanmean(expected_res0, axis=1), axis=0))
+        np.testing.assert_allclose(agg0.cpu(), np.nanmean(np.nanmean(expected_res0, axis=1), axis=0))
         np.testing.assert_equal(not_nans.cpu(), torch.tensor(2))
 
     def test_tolerance_euclidean_distance(self):
@@ -126,7 +126,7 @@ class TestAllSurfaceDiceMetrics(unittest.TestCase):
         expected_res0[1, 1] = np.nan
         for b, c in np.ndindex(batch_size, n_class):
             np.testing.assert_allclose(expected_res0[b, c], res0[b, c].cpu())
-        np.testing.assert_array_equal(agg0.cpu(), np.nanmean(np.nanmean(expected_res0, axis=1), axis=0))
+        np.testing.assert_allclose(agg0.cpu(), np.nanmean(np.nanmean(expected_res0, axis=1), axis=0))
         np.testing.assert_equal(not_nans.cpu(), torch.tensor(2))
 
     def test_tolerance_euclidean_distance_3d(self):
@@ -173,7 +173,7 @@ class TestAllSurfaceDiceMetrics(unittest.TestCase):
         expected_res0[1, 1] = np.nan
         for b, c in np.ndindex(batch_size, n_class):
             np.testing.assert_allclose(expected_res0[b, c], res0[b, c].cpu())
-        np.testing.assert_array_equal(agg0.cpu(), np.nanmean(np.nanmean(expected_res0, axis=1), axis=0))
+        np.testing.assert_allclose(agg0.cpu(), np.nanmean(np.nanmean(expected_res0, axis=1), axis=0))
         np.testing.assert_equal(not_nans.cpu(), torch.tensor(2))
 
     def test_tolerance_all_distances(self):

tests/networks/nets/test_autoencoderkl.py

@@ -330,7 +330,7 @@ class TestAutoEncoderKL(unittest.TestCase):
             weight_path = os.path.join(tmpdir, filename)
             download_url(url=url, filepath=weight_path, hash_val=hash_val, hash_type=hash_type)
 
-            net.load_old_state_dict(torch.load(weight_path), verbose=False)
+            net.load_old_state_dict(torch.load(weight_path, weights_only=True), verbose=False)
 
 
 if __name__ == "__main__":

tests/networks/nets/test_controlnet.py

@@ -208,7 +208,7 @@ class TestControlNet(unittest.TestCase):
             weight_path = os.path.join(tmpdir, filename)
             download_url(url=url, filepath=weight_path, hash_val=hash_val, hash_type=hash_type)
 
-            net.load_old_state_dict(torch.load(weight_path), verbose=False)
+            net.load_old_state_dict(torch.load(weight_path, weights_only=True), verbose=False)
 
 
 if __name__ == "__main__":

tests/networks/nets/test_diffusion_model_unet.py

@@ -578,7 +578,7 @@ class TestDiffusionModelUNet3D(unittest.TestCase):
             weight_path = os.path.join(tmpdir, filename)
             download_url(url=url, filepath=weight_path, hash_val=hash_val, hash_type=hash_type)
 
-            net.load_old_state_dict(torch.load(weight_path), verbose=False)
+            net.load_old_state_dict(torch.load(weight_path, weights_only=True), verbose=False)
 
 
 if __name__ == "__main__":

tests/networks/nets/test_network_consistency.py

@@ -55,7 +55,7 @@ class TestNetworkConsistency(unittest.TestCase):
         print("JSON path: " + json_path)
 
         # Load data
-        loaded_data = torch.load(data_path)
+        loaded_data = torch.load(data_path, weights_only=True)
 
         # Load json from file
         json_file = open(json_path)