map location model to cpu

test_few_shot.py

@@ -31,7 +31,7 @@ def test_main_model(opts):
         st.write("Loading Model Weight...")
     model_main = ModelMain(opts)
     path_ckpt = os.path.join(f"{opts.model_path}")
-    model_main.load_state_dict(torch.load(path_ckpt)['model'])
+    model_main.load_state_dict(torch.load(path_ckpt)['model'], map_location=torch.device('cpu'))
     model_main.to(device)
     model_main.eval()
     with torch.no_grad():

train.py

@@ -35,7 +35,6 @@ def train_main_model(opts):
     val_loader = get_loader(opts.data_root, opts.img_size, opts.language, opts.char_num, opts.max_seq_len, opts.dim_seq, opts.batch_size_val, 'val')
 
     run = wandb.init(project=opts.wandb_project_name, config=opts) # initialize wandb project
-    text_table = wandb.Table(columns=["epoch", "loss", "ref"])
 
     model_main = ModelMain(opts)
     if torch.cuda.is_available() and opts.multi_gpu: