Update app.py

app.py

@@ -108,7 +108,7 @@ class TextGeneration:
             self.model_name_or_path, use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
             pad_token_id=self.tokenizer.eos_token_id, eos_token_id=self.tokenizer.eos_token_id,
             torch_dtype=DTYPE, low_cpu_mem_usage=False if DEVICE == "cpu" else True
-        ).to(device=DEVICE, non_blocking=True)
+        ).to(device=DEVICE, non_blocking=False)
         _ = self.model.eval()
         device_number = -1 if DEVICE == "cpu" else int(DEVICE.split(":")[-1])
         self.generator = pipeline(self.task, model=self.model, tokenizer=self.tokenizer, device=device_number)
@@ -131,7 +131,8 @@ class TextGeneration:
         )[0]["generated_text"]
 
 
-@st.cache(allow_output_mutation=True)
+#@st.cache(allow_output_mutation=True)
+@st.cache(hash_funcs={torch.nn.parameter.Parameter: lambda _: None})
 def load_text_generator():
     generator = TextGeneration()
     generator.load()