Update app.py

app.py

@@ -4,14 +4,7 @@ import random
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 
 def load_model(model_path, dtype):
-    if dtype == "fp32":
-        torch_dtype = torch.float32
-    elif dtype == "fp16":
-        torch_dtype = torch.float16
-    else:
-        raise ValueError("Invalid dtype. Only 'fp32' or 'fp16' are supported.")
-
-    model = T5ForConditionalGeneration.from_pretrained(model_path, torch_dtype=torch_dtype)
+    model = T5ForConditionalGeneration.from_pretrained(model_path, torch_dtype=dtype)
     return model
 
 def generate(
@@ -120,7 +113,7 @@ additional_inputs = [
         info="A starting point to initiate the generation process"
     ),
     gr.Radio(
-        choices=["fp32", "fp16"],
+        choices=[("fp32", torch.float32), ("fp16", torch.float16)],
         value="fp16",
         label="Model Precision",
         info="fp32 is more precised, fp16 is faster and less memory consuming",