Fixed running on CPU

Model/COSMIC/erc_training/predict_epik.py

@@ -42,8 +42,10 @@ def load_model(model_path, args):
 
     if cuda:
         model.cuda()
+        model.load_state_dict(torch.load(model_path))
+    else:
+        model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
 
-    model.load_state_dict(torch.load(model_path))
     model.eval()
 
     return model
@@ -91,7 +93,7 @@ def parse_cosmic_args():
 
     # Parse arguments input into the cosmic model
     parser.add_argument(
-        "--no-cuda", action="store_true", default=False, help="does not use GPU"
+        "--no-cuda", action="store_true", default=True, help="does not use GPU"
     )
     parser.add_argument(
         "--lr", type=float, default=0.0001, metavar="LR", help="learning rate"

Model/COSMIC/feature_extraction/comet/csk_feature_extract.py

@@ -10,10 +10,8 @@ import comet.src.interactive.functions as interactive
 
 
 class CSKFeatureExtractor:
-    def __init__(self, dir="."):
+    def __init__(self, dir=".", device=0):
         super(CSKFeatureExtractor, self).__init__()
-
-        device = "cpu"
         model_file = os.path.join(
             dir, "comet/pretrained_models/atomic_pretrained_model.pickle"
         )

app/views/__init__.py

@@ -9,7 +9,7 @@ EPIK_MODEL_DIR = PATH_TO_COSMIC + "/erc_training"
 sys.path.append(PATH_TO_COSMIC)
 sys.path.append(EXTRACTORS_PATH)
 sys.path.append(EPIK_MODEL_DIR)
-sys.path.append('.')
+sys.path.append(".")
 from Model.COSMIC.feature_extraction.comet.csk_feature_extract import (
     CSKFeatureExtractor,
 )
@@ -23,7 +23,7 @@ roberta = RobertaModel.from_pretrained(
 )
 roberta.eval()
 
-comet = CSKFeatureExtractor(dir=EXTRACTORS_PATH)
+comet = CSKFeatureExtractor(dir=EXTRACTORS_PATH, device="cpu")
 
 cosmic_args = parse_cosmic_args()