removing apex

NVIDIA · Dec 20, 2023 · 59c8be3 · 59c8be3
1 parent ad62db7
commit 59c8be3
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/makani/networks/afnonet_v2.py b/makani/networks/afnonet_v2.py
@@ -22,7 +22,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from apex.normalization import FusedLayerNorm
 import torch.fft
 from torch.nn.modules.container import Sequential
 from torch.utils.checkpoint import checkpoint_sequential
@@ -282,7 +281,7 @@ def _init_weights(self, m):
             # nn.init.normal_(m.weight, std=0.02)
             if m.bias is not None:
                 nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.LayerNorm) or isinstance(m, FusedLayerNorm) or isinstance(m, nn.InstanceNorm3d):
+        elif isinstance(m, nn.LayerNorm) or isinstance(m, nn.InstanceNorm3d):
             nn.init.constant_(m.bias, 0)
             nn.init.constant_(m.weight, 1.0)