Use property instead of function (AI-Hypercomputer#59)

use property instead of function
FanhaiLu1 · Apr 25, 2024 · 2a7894f · 2a7894f
1 parent 781ec9e
commit 2a7894f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jetstream_pt/engine.py b/jetstream_pt/engine.py
@@ -575,7 +575,7 @@ def _load_from_safetensors(self, path):
   # pylint: disable-next=all
   def load_params(self) -> Params:
     # We want to fix this: load from files
-    with jax.default_device(self.colocated_cpus()):
+    with jax.default_device(self.colocated_cpus):
       if self.env.checkpoint_path:
         if self.env.checkpoint_format == "safetensors":
           return self._load_from_safetensors(self.env.checkpoint_path)