pablodawson commited on
Commit
fde5edf
1 Parent(s): 74aacea

Update convert.py

Browse files
Files changed (1) hide show
  1. convert.py +4 -0
convert.py CHANGED
@@ -176,6 +176,8 @@ def check_final_model(model_id: str, folder: str):
176
  input_ids = torch.arange(10).unsqueeze(0)
177
  pixel_values = torch.randn(1, 3, 224, 224)
178
  input_values = torch.arange(1000).float().unsqueeze(0)
 
 
179
  kwargs = {}
180
  if "input_ids" in sig.parameters:
181
  kwargs["input_ids"] = input_ids
@@ -189,6 +191,8 @@ def check_final_model(model_id: str, folder: str):
189
  kwargs["bbox"] = torch.zeros((1, 10, 4)).long()
190
  if "image" in sig.parameters:
191
  kwargs["image"] = pixel_values
 
 
192
 
193
  if torch.cuda.is_available():
194
  pt_model = pt_model.cuda()
 
176
  input_ids = torch.arange(10).unsqueeze(0)
177
  pixel_values = torch.randn(1, 3, 224, 224)
178
  input_values = torch.arange(1000).float().unsqueeze(0)
179
+ task_inputs = torch.randn(1,77)
180
+
181
  kwargs = {}
182
  if "input_ids" in sig.parameters:
183
  kwargs["input_ids"] = input_ids
 
191
  kwargs["bbox"] = torch.zeros((1, 10, 4)).long()
192
  if "image" in sig.parameters:
193
  kwargs["image"] = pixel_values
194
+ if "task_inputs" in sig.parameters:
195
+ kwargs["task_inputs"] = task_inputs
196
 
197
  if torch.cuda.is_available():
198
  pt_model = pt_model.cuda()