add cuda device in app.py

2024-08-23 12:18:48 +02:00 · 2024-08-23 12:18:48 +02:00 · ac442f82ac
commit ac442f82ac
parent 2d91193a98
1 changed files with 3 additions and 0 deletions
--- a/neural_style_app/app.py
+++ b/neural_style_app/app.py
@ -12,6 +12,9 @@ import torch.nn as nn

 app = Flask(__name__)

+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.set_default_device(device)
+
 # Image transformation
 imsize = 512 if torch.cuda.is_available() else 128
 loader = transforms.Compose([