Merge branch 'main' into release-candidate

# Conflicts: # tests/module_tests/test_trainable.py
LarsKue · Dec 6, 2023 · a580910 · a580910
2 parents ed9f82f + 103a3e2
commit a580910
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/lightning_trainable/launcher/fit.py b/src/lightning_trainable/launcher/fit.py
@@ -91,6 +91,11 @@ def main(args=None):
     if num_threads is not None:
         torch.set_num_threads(num_threads)
 
+    # Set matmul precision
+    matmul_precision = hparams.pop("matmul_precision", None)
+    if matmul_precision is not None:
+        torch.set_float32_matmul_precision(matmul_precision)
+
     # Load the model
     if "model" not in hparams:
         model_class_file = checkpoint_file.parent.parent / "model.txt"