make style

Author: SahilCarterr

tests/others/test_ema.py

@@ -60,27 +60,24 @@ def simulate_backprop(self, unet):
         return unet
 
     def test_from_pretrained(self):
-        #Save the model parameters to a temporary directory
+        # Save the model parameters to a temporary directory
         unet, ema_unet = self.get_models()
         with tempfile.TemporaryDirectory() as tmpdir:
             ema_unet.save_pretrained(tmpdir)
 
-            #Load the EMA model from the saved directory
-            loaded_ema_unet = EMAModel.from_pretrained(
-                tmpdir, model_cls=UNet2DConditionModel,foreach=False
-            )
+            # Load the EMA model from the saved directory
+            loaded_ema_unet = EMAModel.from_pretrained(tmpdir, model_cls=UNet2DConditionModel, foreach=False)
 
-        #Check that the shadow parameters of the loaded model match the original EMA model
+        # Check that the shadow parameters of the loaded model match the original EMA model
         for original_param, loaded_param in zip(ema_unet.shadow_params, loaded_ema_unet.shadow_params):
             assert torch.allclose(original_param, loaded_param, atol=1e-4)
 
-        #Verify that the optimization step is also preserved
+        # Verify that the optimization step is also preserved
         assert loaded_ema_unet.optimization_step == ema_unet.optimization_step
 
-        #Check the decay value
+        # Check the decay value
         assert loaded_ema_unet.decay == ema_unet.decay
 
-
     def test_optimization_steps_updated(self):
         unet, ema_unet = self.get_models()
         # Take the first (hypothetical) EMA step.
@@ -215,29 +212,26 @@ def simulate_backprop(self, unet):
             updated_state_dict.update({k: updated_param})
         unet.load_state_dict(updated_state_dict)
         return unet
+
     def test_from_pretrained(self):
-        #Save the model parameters to a temporary directory
+        # Save the model parameters to a temporary directory
         unet, ema_unet = self.get_models()
         with tempfile.TemporaryDirectory() as tmpdir:
             ema_unet.save_pretrained(tmpdir)
 
-            #Load the EMA model from the saved directory
-            loaded_ema_unet = EMAModel.from_pretrained(
-                tmpdir, model_cls=UNet2DConditionModel,foreach=True
-            )
+            # Load the EMA model from the saved directory
+            loaded_ema_unet = EMAModel.from_pretrained(tmpdir, model_cls=UNet2DConditionModel, foreach=True)
 
-        #Check that the shadow parameters of the loaded model match the original EMA model
+        # Check that the shadow parameters of the loaded model match the original EMA model
         for original_param, loaded_param in zip(ema_unet.shadow_params, loaded_ema_unet.shadow_params):
             assert torch.allclose(original_param, loaded_param, atol=1e-4)
 
-        #Verify that the optimization step is also preserved
+        # Verify that the optimization step is also preserved
         assert loaded_ema_unet.optimization_step == ema_unet.optimization_step
 
-        #Check the decay value
+        # Check the decay value
         assert loaded_ema_unet.decay == ema_unet.decay
 
-
-
     def test_optimization_steps_updated(self):
         unet, ema_unet = self.get_models()
         # Take the first (hypothetical) EMA step.