This commit is contained in:
lvmin
2023-08-09 15:07:09 -07:00
parent c56e2c9fd1
commit c104f30137
3 changed files with 4 additions and 2 deletions
+3 -1
View File
@@ -272,7 +272,9 @@ class SiLU(nn.Module):
class GroupNorm32(nn.GroupNorm):
def forward(self, x):
return super().forward(x)
self.weight.to(torch.float32)
self.bias.to(torch.float32)
return super().forward(x.float()).type(x.dtype)
def conv_nd(dims, *args, **kwargs):