Update decorators.py

decorators.py

@@ -7,30 +7,42 @@ import time
 
 class spaces:
     @staticmethod
-    def GPU(duration=100):
+    def GPU(duration=0):
         def decorator(func):
             @functools.wraps(func)
             def wrapper(*args, **kwargs):
-                if not torch.cuda.is_available():
-                    raise RuntimeError("GPU no está disponible.")
+                try:
+                    # Verifica si la GPU está disponible, si no lanza un error controlado
+                    if torch.cuda.is_available():
+                        device = torch.device('cuda')
+                        print("GPU disponible. Ejecutando en GPU.")
+                    else:
+                        raise RuntimeError("No se detecta la GPU. Asegúrate de que los drivers CUDA están instalados y configurados correctamente.")
                 
-                result = [None]
-                exception = [None]
+                    # Pasar el dispositivo a la función como argumento
+                    kwargs['device'] = device
 
-                def target():
-                    try:
-                        result[0] = func(*args, **kwargs)
-                    except Exception as e:
-                        exception[0] = e
+                    result = [None]
+                    exception = [None]
 
-                thread = threading.Thread(target=target)
-                thread.start()
-                thread.join(duration)
+                    def target():
+                        try:
+                            result[0] = func(*args, **kwargs)
+                        except Exception as e:
+                            exception[0] = e
 
-                if thread.is_alive():
-                    raise TimeoutError(f"La ejecución de la función excedió {duration} segundos.")
-                if exception[0]:
-                    raise exception[0]
-                return result[0]
+                    # Ejecutar la función en un hilo separado
+                    thread = threading.Thread(target=target)
+                    thread.start()
+                    thread.join(duration)
+
+                    if thread.is_alive():
+                        raise TimeoutError(f"La ejecución de la función excedió {duration} segundos.")
+                    if exception[0]:
+                        raise exception[0]
+                    return result[0]
+                except RuntimeError as e:
+                    print(f"Error: {str(e)}")
+                    raise e  # Lanzar el error de GPU para que se maneje adecuadamente
             return wrapper
         return decorator