vai0511 commited on
Commit
545643a
·
verified ·
1 Parent(s): a82fdd3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -3,6 +3,9 @@ from transformers import AutoTokenizer, T5ForConditionalGeneration
3
  import torch
4
  import os
5
 
 
 
 
6
  #access_token = os.getenv["HF_TOKEN"]
7
 
8
  # Load model and tokenizer
 
3
  import torch
4
  import os
5
 
6
+ # Set up device
7
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8
+
9
  #access_token = os.getenv["HF_TOKEN"]
10
 
11
  # Load model and tokenizer