ClearML/clearml
1
0
Fork 0
mirror of https://github.com/clearml/clearml synced 2025-03-03 10:42:00 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix pytorch automagic logging torchscript models

Browse Source
This commit is contained in:
allegroai 2022-04-13 14:16:12 +03:00
parent 3996cdbc7e
commit e06d3cbf99
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
examples/frameworks/pytorch/pytorch_tensorboard.py
View File

@ -131,7 +131,9 @@ def main():
for epoch in range(1, args.epochs + 1):
train(model, epoch, train_loader, args, optimizer, writer)
torch.save(model, os.path.join(gettempdir(), 'model{}'.format(epoch)))
m = torch.jit.script(model)
m.save(os.path.join(gettempdir(), 'model{}'.format(epoch)))
#torch.save(model, os.path.join(gettempdir(), 'model{}'.format(epoch)))
test(model, test_loader, args, optimizer, writer)