Closed ver217 closed 2 weeks ago
from tensornvme.async_file_io import AsyncFileWriter import torch.nn as nn import torch model = nn.Linear(10, 2) with open("t.pkl", "wb") as f: f_writer.synchronize() torch.save(model.state_dict(), f_writer) f_writer.synchronize() print(model.state_dict()) # output # OrderedDict([('weight', # tensor([[-0.2586, -0.0774, -0.1293, 0.0518, 0.1004, 0.0637, -0.1227, 0.0488, # -0.2836, -0.1508], # [ 0.1144, -0.2542, -0.0783, 0.2590, -0.0972, -0.0133, -0.1492, 0.2579, # -0.2547, 0.1043]])), # ('bias', tensor([0.0223, 0.1377]))]) print(torch.load("t.pkl")) # output # OrderedDict([('weight', # tensor([[-0.2586, -0.0774, -0.1293, 0.0518, 0.1004, 0.0637, -0.1227, 0.0488, # -0.2836, -0.1508], # [ 0.1144, -0.2542, -0.0783, 0.2590, -0.0972, -0.0133, -0.1492, 0.2579, # -0.2547, 0.1043]])), # ('bias', tensor([0.0223, 0.1377]))])
Usage