bertugmirasyedi commited on
Commit
34183cc
·
1 Parent(s): dc12110

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -0
app.py CHANGED
@@ -8,6 +8,7 @@ from transformers import (
8
  )
9
  from optimum.onnxruntime import ORTModelForSeq2SeqLM, ORTModelForSequenceClassification
10
  from sentence_transformers import SentenceTransformer
 
11
 
12
  # Define the FastAPI app
13
  app = FastAPI(docs_url="/")
@@ -21,6 +22,10 @@ app.add_middleware(
21
  allow_headers=["*"],
22
  )
23
 
 
 
 
 
24
  # Define the Google Books API key
25
  key = os.environ.get("GOOGLE_BOOKS_API_KEY")
26
 
 
8
  )
9
  from optimum.onnxruntime import ORTModelForSeq2SeqLM, ORTModelForSequenceClassification
10
  from sentence_transformers import SentenceTransformer
11
+ import torch
12
 
13
  # Define the FastAPI app
14
  app = FastAPI(docs_url="/")
 
22
  allow_headers=["*"],
23
  )
24
 
25
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
26
+ # True
27
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
28
+
29
  # Define the Google Books API key
30
  key = os.environ.get("GOOGLE_BOOKS_API_KEY")
31