cpu support added

FGT_codes/tool/video_inpainting.py

@@ -181,7 +181,10 @@ def create_dir(dir):
 def initialize_RAFT(args, device):
     """Initializes the RAFT model."""
     model = torch.nn.DataParallel(RAFT(args))
-    model.load_state_dict(torch.load(args.raft_model))
+    if not torch.cuda.is_available():
+        model.load_state_dict(torch.load(args.raft_model, map_location=lambda storage, loc: storage))
+    else:
+        model.load_state_dict(torch.load(args.raft_model))
 
     model = model.module
     model.to(device)
@@ -202,9 +205,14 @@ def initialize_LAFC(args, device):
     model = configs["model"]
     pkg = import_module("LAFC.models.{}".format(model))
     model = pkg.Model(configs)
-    state = torch.load(
-        checkpoint, map_location=lambda storage, loc: storage.cuda(device)
-    )
+    if not torch.cuda.is_available():
+        state = torch.load(
+        checkpoint, map_location=lambda storage, loc: storage
+        )
+    else:
+        state = torch.load(
+            checkpoint, map_location=lambda storage, loc: storage.cuda(device)
+        )
     model.load_state_dict(state["model_state_dict"])
     model = model.to(device)
     return model, configs
@@ -221,9 +229,14 @@ def initialize_FGT(args, device):
     model = configs["model"]
     net = import_module("FGT.models.{}".format(model))
     model = net.Model(configs).to(device)
-    state = torch.load(
-        checkpoint, map_location=lambda storage, loc: storage.cuda(device)
-    )
+    if not torch.cuda.is_available():
+        state = torch.load(
+        checkpoint, map_location=lambda storage, loc: storage
+        )
+    else:
+        state = torch.load(
+            checkpoint, map_location=lambda storage, loc: storage.cuda(device)
+        )
     model.load_state_dict(state["model_state_dict"])
     return model, configs