Update model/model.py

model/model.py

@@ -6,6 +6,8 @@ from transformers import AutoModelForImageClassification
 def predict(image_path):
     model = AutoModelForImageClassification.from_pretrained('sensei-ml/concrete_crack_images_classification')
     model.eval()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
 
     with torch.no_grad():
         # Convertir el array de NumPy a un tensor de PyTorch
@@ -13,7 +15,7 @@ def predict(image_path):
 
         # Redimensionar la imagen usando funciones de transformación que soporten tensores
         image_tensor = F.resize(image_tensor, [227, 227])
-
+       
         # Normalización
         transform = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 
@@ -21,6 +23,7 @@ def predict(image_path):
         input_tensor = transform(image_tensor).unsqueeze(0)  # Añadir la dimensión del batch
 
         # Hacer predicciones
+        input_tensor = input_tensor.to(device)
         outputs = model(input_tensor)
         logits = outputs.logits
         probabilities = torch.nn.functional.softmax(logits, dim=1)[0]