From ac442f82ac5eba8a589e31e3bffc59459da3e4c6 Mon Sep 17 00:00:00 2001
From: s486797 <Sarafina1*>
Date: Fri, 23 Aug 2024 12:18:48 +0200
Subject: [PATCH] add cuda device in app.py

---
 neural_style_app/app.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/neural_style_app/app.py b/neural_style_app/app.py
index 4b3b7e9..d235aea 100644
--- 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([