Fix tests formatting.

Author: Virginia Fernandez

tests/inferers/test_diffusion_inferer.py

@@ -100,12 +100,7 @@ def test_sample_cfg(self, model_params, input_shape):
         inferer = DiffusionInferer(scheduler=scheduler)
         scheduler.set_timesteps(num_inference_steps=10)
         sample, intermediates = inferer.sample(
-            input_noise=noise,
-            network=model,
-            scheduler=scheduler,
-            save_intermediates=True,
-            intermediate_steps=1,
-            cfg=5,
+            input_noise=noise, network=model, scheduler=scheduler, save_intermediates=True, intermediate_steps=1, cfg=5
         )
         self.assertEqual(sample.shape, noise.shape)
 

tests/inferers/test_latent_diffusion_inferer.py

@@ -402,11 +402,7 @@ def test_sample_shape(
                     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,
-                    network=stage_2,
-                    scheduler=scheduler,
-                    seg=input_seg,
+                    input_noise=noise, autoencoder_model=stage_1, network=stage_2, scheduler=scheduler, seg=input_seg
                 )
             else:
                 sample = inferer.sample(
@@ -567,11 +563,7 @@ def test_get_likelihoods(
             )
         else:
             sample, intermediates = inferer.get_likelihood(
-                inputs=input,
-                autoencoder_model=stage_1,
-                network=stage_2,
-                scheduler=scheduler,
-                save_intermediates=True,
+                inputs=input, autoencoder_model=stage_1, network=stage_2, scheduler=scheduler, save_intermediates=True
             )
         self.assertEqual(len(intermediates), 10)
         self.assertEqual(intermediates[0].shape, latent_shape)
@@ -927,11 +919,7 @@ def test_incompatible_spade_setup(self):
 
             with self.assertRaises(ValueError):
                 _ = inferer.sample(
-                    input_noise=noise,
-                    autoencoder_model=stage_1,
-                    network=stage_2,
-                    scheduler=scheduler,
-                    seg=input_seg,
+                    input_noise=noise, autoencoder_model=stage_1, network=stage_2, scheduler=scheduler, seg=input_seg
                 )