use cuda if available

train_semeval/__main__.py

@@ -97,6 +97,8 @@ if __name__ == "__main__":
     emb_dim = 64
     voc_size = len(voc)
     model = MiniBertForTraining(emb_dim, voc_size, mask_idx, hidden_dim=64)
+    if torch.cuda.is_available():
+        model = model.to("cuda")
     learning_rate = 1e-3
     optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
 

train_semeval/init_with_w2v.py

@@ -109,7 +109,8 @@ if __name__ == "__main__":
 
     model = MiniBertForTraining(emb_dim, voc_size, mask_idx, hidden_dim=64)
     model.minibert.embedding.word_embeddings.weight = torch.nn.Parameter(embs)
-
+    if torch.cuda.is_available():
+        model = model.to("cuda")
     learning_rate = 1e-3
     optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)