Seokju Cho commited on
Commit
5d50721
·
1 Parent(s): 9e9e6ff
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -132,9 +132,9 @@ def extract_feature(video_input, model_size="small"):
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
 
 
132
  with torch.no_grad():
133
  feature = model.get_feature_grids(video_input)
134
 
135
+ feature.lowres = (x.cpu() for x in feature.lowres)
136
+ feature.hires = (x.cpu() for x in feature.hires)
137
+ feature.highest = (x.cpu() for x in feature.highest)
138
 
139
  return feature
140