Merge pull request #46 from loopinf/feature/32-fix_bug_F_linear

Feature/32 fix bug f linear
EleutherAI · Sep 28, 2022 · 7af33eb · 7af33eb
2 parents 88ac525 + cf976a1
commit 7af33eb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/oslo/torch/nn/modules/linear.py b/oslo/torch/nn/modules/linear.py
@@ -17,12 +17,13 @@ def __init__(
         bias: bool = True,
         dtype: Optional[torch.dtype] = None,
         skip_bias_add: bool = False,
+        device: torch.device = None,
     ):
         super(Linear, self).__init__(
             in_features=in_features,
             out_features=out_features,
             bias=bias,
-            device=torch.device(torch.cuda.current_device()),
+            device=device,
             dtype=dtype,
         )
         self.skip_bias_add = skip_bias_add