Commit 41e47246 authored by Gaëtan Caillaut's avatar Gaëtan Caillaut
Browse files

models can be saved

parent 99bf74f0
......@@ -41,6 +41,7 @@ if __name__ == "__main__":
help="Use GPU acceleration")
parser.add_argument("--activation", default="gelu", type=str)
parser.add_argument("--logdir", default=None, type=str)
parser.add_argument("--save", type=str, required=False)
args = parser.parse_args()
......@@ -127,7 +128,8 @@ if __name__ == "__main__":
writer.flush()
writer.close()
model.set_train(False)
if args.save is not None:
torch.save(model.state_dict(), args.save)
# test_sentences = crps[:3]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment