🎓 A1Learner
🔍

PyTorch Saving Tensors

torch.save writes a tensor into a file on disk. torch.load reads the file back and returns the same tensor.

Save a tensor to a file

torch.save writes the tensor into a file with the given name. The file lands in the folder the program runs from.

import torch

table = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
torch.save(table, "table.pt")
print("saved")

Load the tensor back

torch.load reads the file and returns the tensor that was saved.

loaded = torch.load("table.pt")
print(loaded)

Check the two tensors match

torch.equal compares shape and values. The loaded tensor matches the saved one.

print(torch.equal(table, loaded))

Example

import torch

table = torch.tensor([[10.0, 20.0],
                      [30.0, 40.0]])
torch.save(table, "table.pt")
loaded = torch.load("table.pt")
print("Saved table:")
print(table)
print("Loaded table:")
print(loaded)
print("Equal:", torch.equal(table, loaded))

Key points

  • Use torch.save(tensor, "name.pt") to write a tensor to a file
  • Use torch.load("name.pt") to read a saved tensor back
  • The file lands in the folder the program runs from
  • Use torch.equal(a, b) to check two tensors match in shape and values
  • A saved tensor keeps its values, shape, and dtype
  • Saving turns a tensor into a file, and loading turns the file back into a tensor

In this example, table holds four values in two rows. torch.save writes the table into a file named table.pt. The file lands in the folder this program runs from. torch.load reads the file and returns the saved tensor. The saved and loaded tables print the same. torch.equal confirms both tensors match, so it prints True.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Where does torch.save put the file when the name has no folder part?