Update instance_norm.py

segyges · Apr 19, 2023 · 324cee0 · 324cee0
1 parent 988cafe
commit 324cee0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/apex/normalization/instance_norm.py b/apex/normalization/instance_norm.py
@@ -129,6 +129,8 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
     def forward(self, input: Tensor) -> Tensor:
         assert input.is_cuda, "NVFuser InstanceNorm is CUDA only"
         self._check_input_dim(input)
+        if self.dummy.device != input.device:
+            self.dummy = torch.empty([], device=input.device)
         if self.running_mean is not None:
             out = InstanceNormNVFuserFunction.apply(
                     input, self.weight if self.weight is not None else self.dummy,