Fix type hint in initialize_torch_settings function

RichieHakim · Apr 6, 2024 · 5fd5db6 · 5fd5db6
1 parent 7f2a0aa
commit 5fd5db6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bnpm/torch_helpers.py b/bnpm/torch_helpers.py
@@ -224,7 +224,7 @@ def initialize_torch_settings(
     enable_cudnn: Optional[bool] = None,
     deterministic_cudnn: Optional[bool] = None,
     deterministic_torch: Optional[bool] = None,
-    set_global_device: Optional[str, torch.device] = None,
+    set_global_device: Optional[Union[str, torch.device]] = None,
     init_linalg: bool = True,
     init_linalg_device: Union[str, torch.device] = 'cuda:0',
 ) -> None: