use kwargs

Author: SkafteNicki

examples/fabric/image_classifier/train_fabric.py

@@ -158,7 +158,7 @@ def run(hparams):
     # When using distributed training, use `fabric.save`
     # to ensure the current process is allowed to save a checkpoint
     if hparams.save_model:
-        fabric.save("mnist_cnn.pt", model.state_dict())
+        fabric.save(path="mnist_cnn.pt", state=model.state_dict())
 
 
 if __name__ == "__main__":

examples/fabric/kfold_cv/train_fabric.py

@@ -161,7 +161,7 @@ def run(hparams):
     # When using distributed training, use `fabric.save`
     # to ensure the current process is allowed to save a checkpoint
     if hparams.save_model:
-        fabric.save("mnist_cnn.pt", model.state_dict())
+        fabric.save(path="mnist_cnn.pt", state=model.state_dict())
 
 
 if __name__ == "__main__":

examples/fabric/tensor_parallel/train.py

@@ -67,7 +67,7 @@ def train():
     # See `fabric consolidate --help` if you need to convert the checkpoint to a single file
     fabric.print("Saving a (distributed) checkpoint ...")
     state = {"model": model, "optimizer": optimizer, "iteration": i}
-    fabric.save("checkpoint.pt", state)
+    fabric.save(path="checkpoint.pt", state=state)
 
     fabric.print("Training successfully completed!")
     fabric.print(f"Peak memory usage: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB")