style

Author: sayakpaul

tests/models/autoencoders/test_models_autoencoder_oobleck.py

@@ -114,7 +114,9 @@ def test_forward_with_norm_groups(self):
     def test_set_attn_processor_for_determinism(self):
         return
 
-    @unittest.skip("Test not supported because of 'weight_norm_fwd_first_dim_kernel' not implemented for 'Float8_e4m3fn'")
+    @unittest.skip(
+        "Test not supported because of 'weight_norm_fwd_first_dim_kernel' not implemented for 'Float8_e4m3fn'"
+    )
     def test_layerwise_casting_training(self):
         return super().test_layerwise_casting_training()
 

tests/models/test_modeling_common.py

@@ -1338,7 +1338,7 @@ def test_variant_sharded_ckpt_right_format(self):
                 # Example: diffusion_pytorch_model.fp16-00001-of-00002.safetensors
                 assert all(f.split(".")[1].split("-")[0] == variant for f in shard_files)
 
-    def test_layerwise_casting_training(self):    
+    def test_layerwise_casting_training(self):
         def test_fn(storage_dtype, compute_dtype):
             if torch.device(torch_device).type == "cpu" and compute_dtype == torch.bfloat16:
                 return
@@ -1360,15 +1360,14 @@ def test_fn(storage_dtype, compute_dtype):
                 noise = torch.randn((input_tensor.shape[0],) + self.output_shape).to(torch_device)
                 noise = cast_maybe_tensor_dtype(noise, torch.float32, compute_dtype)
                 loss = torch.nn.functional.mse_loss(output, noise)
-            
+
             loss.backward()
 
         test_fn(torch.float16, torch.float32)
         test_fn(torch.float8_e4m3fn, torch.float32)
         test_fn(torch.float8_e5m2, torch.float32)
         test_fn(torch.float8_e4m3fn, torch.bfloat16)
 
-
     def test_layerwise_casting_inference(self):
         from diffusers.hooks.layerwise_casting import DEFAULT_SKIP_MODULES_PATTERN, SUPPORTED_PYTORCH_LAYERS