Initialize default device lazily (#152)

This way, we can import the library without requiring a GPU to be present.
databricks · Sep 19, 2024 · 7b0337f · 7b0337f
1 parent cc7614e
commit 7b0337f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megablocks/layers/arguments.py b/megablocks/layers/arguments.py
@@ -52,7 +52,7 @@ class Arguments:
     # Initialization arguments.
     fp16: bool = True
     bf16: bool = False
-    device: Union[int, torch.device] = torch.cuda.current_device()
+    device: Union[int, torch.device] = dataclasses.field(default_factory=torch.cuda.current_device)
     init_method: InitFn = partial(torch.nn.init.normal_, mean=0.0, std=0.02)
     output_layer_init_method: InitFn = init_method