fix GroupNorm init

Oneflow-Inc · Apr 19, 2024 · a42da45 · a42da45
1 parent 44ad994
commit a42da45
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/oneflow/nn/modules/normalization.py b/python/oneflow/nn/modules/normalization.py
@@ -137,8 +137,8 @@ def __init__(
         if dtype:
             factory_kwargs["dtype"] = dtype
         if self.affine:
-            self.weight = flow.nn.Parameter(flow.Tensor(num_channels, **factory_kwargs))
-            self.bias = flow.nn.Parameter(flow.Tensor(num_channels, **factory_kwargs))
+            self.weight = flow.nn.Parameter(flow.Tensor(num_channels).to(**factory_kwargs))
+            self.bias = flow.nn.Parameter(flow.Tensor(num_channels).to(**factory_kwargs))
         else:
             self.register_parameter("weight", None)
             self.register_parameter("bias", None)