sindhuhegde commited on
Commit
43bd4b0
1 Parent(s): e41c892

Update app

Browse files
Files changed (1) hide show
  1. app.py +4 -5
app.py CHANGED
@@ -34,6 +34,7 @@ CHECKPOINT_PATH = "model_rgb.pth"
34
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
35
  use_cuda = torch.cuda.is_available()
36
  n_negative_samples = 100
 
37
 
38
  def preprocess_video(path, result_folder, apply_preprocess, padding=20):
39
 
@@ -225,7 +226,9 @@ def load_checkpoint(path, model):
225
  for k, v in s.items():
226
  new_s[k.replace('module.', '')] = v
227
  model.load_state_dict(new_s)
228
- model.cuda()
 
 
229
 
230
  print("Loaded checkpoint from: {}".format(path))
231
 
@@ -721,10 +724,6 @@ def process_video(video_path, num_avg_frames, apply_preprocess):
721
 
722
  # Convert frames to tensor
723
  rgb_frames = np.transpose(rgb_frames, (4, 0, 1, 2, 3))
724
- print("Transposed array: ", rgb_frames.shape)
725
- a = np.array([2,2])
726
- b = torch.tensor(a)
727
- print("checking torch: ", b.shape)
728
  rgb_frames = torch.FloatTensor(rgb_frames).unsqueeze(0)
729
  B = rgb_frames.size(0)
730
  print("Successfully converted the frames to tensor")
 
34
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
35
  use_cuda = torch.cuda.is_available()
36
  n_negative_samples = 100
37
+ print("Using CUDA: ", use_cuda, device)
38
 
39
  def preprocess_video(path, result_folder, apply_preprocess, padding=20):
40
 
 
226
  for k, v in s.items():
227
  new_s[k.replace('module.', '')] = v
228
  model.load_state_dict(new_s)
229
+
230
+ if use_cuda:
231
+ model.cuda()
232
 
233
  print("Loaded checkpoint from: {}".format(path))
234
 
 
724
 
725
  # Convert frames to tensor
726
  rgb_frames = np.transpose(rgb_frames, (4, 0, 1, 2, 3))
 
 
 
 
727
  rgb_frames = torch.FloatTensor(rgb_frames).unsqueeze(0)
728
  B = rgb_frames.size(0)
729
  print("Successfully converted the frames to tensor")