updates

Author: sayakpaul

tests/models/autoencoders/test_models_autoencoder_oobleck.py

@@ -114,6 +114,10 @@ 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'")
+    def test_layerwise_casting_training(self):
+        return super().test_layerwise_casting_training()
+
     @unittest.skip(
         "The convolution layers of AutoencoderOobleck are wrapped with torch.nn.utils.weight_norm. This causes the hook's pre_forward to not "
         "cast the module weights to compute_dtype (as required by forward pass). As a result, forward pass errors out. To fix:\n"

tests/models/test_modeling_common.py

@@ -1350,18 +1350,19 @@ def test_fn(storage_dtype, compute_dtype):
             model.train()
 
             inputs_dict = cast_maybe_tensor_dtype(inputs_dict, torch.float32, compute_dtype)
-            output = model(**inputs_dict)
+            with torch.amp.autocast(device_type=torch.device(torch_device).type):
+                output = model(**inputs_dict)
 
-            if isinstance(output, dict):
-                output = output.to_tuple()[0]
+                if isinstance(output, dict):
+                    output = output.to_tuple()[0]
 
-            input_tensor = inputs_dict[self.main_input_name]
-            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)
+                input_tensor = inputs_dict[self.main_input_name]
+                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)

tests/models/unets/test_models_unet_1d.py

@@ -60,6 +60,10 @@ def test_ema_training(self):
     def test_training(self):
         pass
 
+    @unittest.skip("Test not supported.")
+    def test_layerwise_casting_training(self):
+        pass
+
     def test_determinism(self):
         super().test_determinism()
 
@@ -239,6 +243,10 @@ def test_ema_training(self):
     def test_training(self):
         pass
 
+    @unittest.skip("Test not supported.")
+    def test_layerwise_casting_training(self):
+        pass
+
     def prepare_init_args_and_inputs_for_common(self):
         init_dict = {
             "in_channels": 14,