Update app.py

app.py

@@ -10,6 +10,8 @@ from VAE_model_tablets_class import VAE
 
 import gradio as gr
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 num_classes = len(TabletPeriodDataset.PERIOD_INDICES)
 
 class_weights = torch.load("class_weights_period.pt")