Update Garage/AugmenterPipeline.py

Garage/AugmenterPipeline.py

@@ -97,6 +97,8 @@ class Augmenter:
         if mask.mode != 'L': 
             mask = mask.convert('L')
 
+        self.to("cuda")
+
         image_description = self._models["LLaVA"].generate_image_description(image)
         prompt, new_object = self._models["LLaMA"].generate_prompt(current_object, image_description, new_objects_list)
         # prompt = "cat"