From e06d3cbf99e34da8606375341e7826b106aac881 Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Wed, 13 Apr 2022 14:16:12 +0300 Subject: [PATCH] Fix pytorch automagic logging torchscript models --- examples/frameworks/pytorch/pytorch_tensorboard.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/frameworks/pytorch/pytorch_tensorboard.py b/examples/frameworks/pytorch/pytorch_tensorboard.py index cd3cbfa8..b7016cc9 100644 --- a/examples/frameworks/pytorch/pytorch_tensorboard.py +++ b/examples/frameworks/pytorch/pytorch_tensorboard.py @@ -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)