Skip to main content

/

Docs

Nightly

v24.6

Log in

Mojo function

save

save[PathLike: PathLike](tensor_dict: TensorDict, path: PathLike) -> object

Saves a collection of tensors to a file.

The file is saved in a binary format that's specific to MAX. You can then load the checkpoint with load().

For example:

from max.graph.checkpoint import save, TensorDict
from max.tensor import Tensor, TensorShape

def write_to_disk():
    tensors = TensorDict()
    tensors.set("x", Tensor[DType.int32](TensorShape(1, 2, 2), 1, 2, 3, 4))
    tensors.set("y", Tensor[DType.float32](TensorShape(10, 5), -1.23))
    save(tensors, "/path/to/checkpoint.maxckpt")
from max.graph.checkpoint import save, TensorDict
from max.tensor import Tensor, TensorShape

def write_to_disk():
    tensors = TensorDict()
    tensors.set("x", Tensor[DType.int32](TensorShape(1, 2, 2), 1, 2, 3, 4))
    tensors.set("y", Tensor[DType.float32](TensorShape(10, 5), -1.23))
    save(tensors, "/path/to/checkpoint.maxckpt")

Args:

  • ​tensor_dict (TensorDict): Tensors to save.
  • ​path (PathLike): The location to save the checkpoint file. You can use whatever filename and file extension you want.