Removing .float() (autocast in fp16 will discard this (I think)). (…

…open-mmlab#495)
aravind-h-v · Sep 14, 2022 · 7c4b38b · 7c4b38b
1 parent ab7a78e
commit 7c4b38b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/diffusers/models/resnet.py b/src/diffusers/models/resnet.py
@@ -333,7 +333,7 @@ def forward(self, x, temb):
 
         # make sure hidden states is in float32
         # when running in half-precision
-        hidden_states = self.norm1(hidden_states.float()).type(hidden_states.dtype)
+        hidden_states = self.norm1(hidden_states).type(hidden_states.dtype)
         hidden_states = self.nonlinearity(hidden_states)
 
         if self.upsample is not None:
@@ -351,7 +351,7 @@ def forward(self, x, temb):
 
         # make sure hidden states is in float32
         # when running in half-precision
-        hidden_states = self.norm2(hidden_states.float()).type(hidden_states.dtype)
+        hidden_states = self.norm2(hidden_states).type(hidden_states.dtype)
         hidden_states = self.nonlinearity(hidden_states)
 
         hidden_states = self.dropout(hidden_states)