Seokju Cho commited on
Commit
9e9e6ff
1 Parent(s): 72bbdf9
Files changed (1) hide show
  1. app.py +4 -0
app.py CHANGED
@@ -131,6 +131,10 @@ def extract_feature(video_input, model_size="small"):
131
  with torch.autocast(device_type=device, dtype=dtype):
132
  with torch.no_grad():
133
  feature = model.get_feature_grids(video_input)
 
 
 
 
134
 
135
  return feature
136
 
 
131
  with torch.autocast(device_type=device, dtype=dtype):
132
  with torch.no_grad():
133
  feature = model.get_feature_grids(video_input)
134
+
135
+ feature.lowres = feature.lowres.cpu()
136
+ feature.hires = feature.hires.cpu()
137
+ feature.highest = feature.highest.cpu()
138
 
139
  return feature
140