pablodawson
commited on
Commit
•
fde5edf
1
Parent(s):
74aacea
Update convert.py
Browse files- 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()
|