import torch
from safetensors.torch import load_file, save_file

tensors = {'X': torch.randn(1,2,3,4,5,6,7,8,9),}
print('shape:', tensors['X'].shape)
save_file(tensors, 'test_torch.safetensors')

loaded = load_file('test_torch.safetensors')
assert loaded['X'].shape == (1,2,3,4,5,6,7,8,9)
print('loaded:', loaded['X'].shape)

import os
os.remove('test_torch.safetensors')
