s486797/neural_style
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add cuda device in app.py

Browse Source
This commit is contained in:
s486797 2024-08-23 12:18:48 +02:00
parent 2d91193a98
commit ac442f82ac
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
neural_style_app/app.py
View File

@ -12,6 +12,9 @@ import torch.nn as nn
app = Flask(__name__) app = Flask(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_default_device(device)
# Image transformation # Image transformation
imsize = 512 if torch.cuda.is_available() else 128 imsize = 512 if torch.cuda.is_available() else 128
loader = transforms.Compose([ loader = transforms.Compose([