!!

app.py

@@ -150,7 +150,7 @@ model = torch.load('complete-model.pt', map_location=device)
 # inference
 slider_value = st.slider('Amount of text to generate', min_value=100, max_value=2000, value=500, step=5)
 if st.button('Generat text'):
-    context = torch.zeros((1, 1), dtype=torch.long, device='cuda')
+    context = torch.zeros((1, 1), dtype=torch.long, device=device)
     text = model.generate(context, max_new_tokens=slider_value)[0].tolist()
     st.json(decode(text))
 #