Update detect_from_videos.py

detect_from_videos.py

@@ -11,6 +11,9 @@ from tqdm import tqdm
 from model_core import Two_Stream_Net
 from torchvision import transforms
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
 xception_default_data_transforms_256 = {
     'train': transforms.Compose([
         transforms.Resize((256, 256)),
@@ -148,6 +151,7 @@ def test_full_image_network(video_path, model_path, output_path,
     # model, *_ = model_selection(modelname='xception', num_out_classes=2)
     model = Two_Stream_Net()
     model.load_state_dict(torch.load(model_path))
+    model = model.to(device)
     model.eval()
 
     if cuda: