minor changes to app.py
Browse files
app.py
CHANGED
|
@@ -14,7 +14,7 @@ import torch
|
|
| 14 |
physical_devices = tf.config.experimental.list_physical_devices('GPU')
|
| 15 |
if len(physical_devices) > 0:
|
| 16 |
tf.config.experimental.set_memory_growth(physical_devices[0], True)
|
| 17 |
-
device = "cpu"
|
| 18 |
|
| 19 |
@st.cache_resource
|
| 20 |
def load_text_model():
|
|
|
|
| 14 |
physical_devices = tf.config.experimental.list_physical_devices('GPU')
|
| 15 |
if len(physical_devices) > 0:
|
| 16 |
tf.config.experimental.set_memory_growth(physical_devices[0], True)
|
| 17 |
+
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
| 18 |
|
| 19 |
@st.cache_resource
|
| 20 |
def load_text_model():
|