Move to cuda

app.py

@@ -4,7 +4,7 @@ import spaces
 from standalone_velvet import setup_models
 
 models_dict = setup_models("visual_bloom.torch")
-visual_bloom = models_dict["visual_bloom"]
+visual_bloom = models_dict["visual_bloom"].to('cuda')
 tokenizer = models_dict["tokenizer"]
 image_feature_collator = models_dict["image_feature_collator"]
 
@@ -14,10 +14,10 @@ def run_inference(text_input, image_input):
     image_features, image_attentions = image_feature_collator([image_input])
     instruction_inputs = tokenizer([text_input], return_tensors="pt")
     language_output = visual_bloom.generate(
-        image_features,
-        image_attentions,
-        instruction_inputs["input_ids"],
-        instruction_inputs["attention_mask"],
+        image_features.to('cuda'),
+        image_attentions.to('cuda'),
+        instruction_inputs["input_ids"].to('cuda'),
+        instruction_inputs["attention_mask"].to('cuda'),
     )
 
     human_output = tokenizer.decode(language_output[0], skip_special_tokens=True)