Added dummy input for one-shot size-checking

2024-02-29 18:44:42 +00:00 · 2024-02-29 18:44:42 +00:00 · ab633da4a5
commit ab633da4a5
parent dcc3f57596
4 changed files with 21 additions and 7 deletions
--- a/checkpoint.py
+++ b/checkpoint.py
@ -0,0 +1,8 @@
+import torch
+import shutil
+
+def save_checkpoint(state, is_best: bool, task_id, fname="checkpoint.pth.tar"):
+    fdir = "./"+str(task_id)+"/"
+    torch.save(state, fdir + fname)
+    if is_best:
+        shutil.copyfile(fdir + fname, fdir + "best.pth.tar")