fix cuda

app.py

@@ -7,6 +7,7 @@ from st_utils import (
 )
 from huggingface_hub import hf_hub_download
 import os
+import torch
 
 # list_files(os.getcwd())
 
@@ -19,6 +20,8 @@ Simply select one of the sample Python functions from the dropdown menu below, a
 """
 )
 
+st.write(f"Has CUDA: {torch.cuda.is_available()}")
+
 # Download the model from the Hugging Face Hub if it doesn't exist
 download_model()
 

model.py

@@ -104,7 +104,10 @@ class Seq2Seq(nn.Module):
         else:
             # Predict
             preds = []
-            zero = torch.cuda.LongTensor(1).fill_(0)
+            try:
+                zero = torch.cuda.LongTensor(1).fill_(0)
+            except Exception as e:
+                zero = torch.LongTensor(1).fill_(0)
             for i in range(source_ids.shape[0]):
                 context = encoder_output[:, i : i + 1]
                 context_mask = source_mask[i : i + 1, :]
@@ -154,7 +157,10 @@ class Seq2Seq(nn.Module):
 class Beam(object):
     def __init__(self, size, sos, eos):
         self.size = size
-        self.tt = torch.cuda
+        if torch.cuda.is_available():
+            self.tt = torch.cuda
+        else:
+            self.tt = torch
         # The score for each translation on the beam.
         self.scores = self.tt.FloatTensor(size).zero_()
         # The backpointers at each time-step.