Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces # for using hugging face ZeroGPU
 import gradio as gr
 import torch
 from torch import nn
@@ -88,7 +89,8 @@ def load_checkpoint_for_inference(filepath, model_class):
     # loaded_model = load_checkpoint_for_inference("model_path", ClassConditionedUnet)
 
 
-
+# Make this function run on GPU
+@spaces.GPU
 def generate_images(selected_class_name, num_samples):
     print(f"Generating {num_samples} samples for class: {selected_class_name}")