fix device for dummy tensor

segyges · Apr 19, 2023 · 62a2ff9 · 62a2ff9
1 parent d94d55a
commit 62a2ff9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/apex/normalization/instance_norm.py b/apex/normalization/instance_norm.py
@@ -91,7 +91,7 @@ def __init__(
         factory_kwargs = {'device': device, 'dtype': dtype}
         super(_InstanceNormNVFuser, self).__init__(
             num_features, eps, momentum, affine, track_running_stats, **factory_kwargs)
-        self.dummy = torch.empty([], device='cuda')
+        self.dummy = torch.empty([], device=device)
 
     def _check_input_dim(self, input):
         raise NotImplementedError