Removing `.float()` (`autocast` in fp16 will discard this (I think)). (#495)

7c4b38ba · Nicolas Patry · GitHub · ab7a78e8 · 7c4b38ba
Unverified Commit 7c4b38ba authored Sep 14, 2022 by Nicolas Patry Committed by GitHub Sep 14, 2022
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/diffusers/models/resnet.py src/diffusers/models/resnet.py +2 -2

No files found.
--- a/src/diffusers/models/resnet.py
+++ b/src/diffusers/models/resnet.py
@@ -333,7 +333,7 @@ class ResnetBlock2D(nn.Module):

        # 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 @@ class ResnetBlock2D(nn.Module):

        # 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)