torch.load and torch.save
./code/load-and-save.py
1#!/usr/bin/env python3
2
3import torch
4import tempfile
5
6
7def main():
8 a = torch.arange(3)
9 with tempfile.NamedTemporaryFile() as f:
10 torch.save(a, f)
11 f.seek(0)
12 b = torch.load(f)
13 assert torch.all(torch.eq(a, b)), (a, b)
14
15
16if __name__ == "__main__":
17 main()