From 34bfe78218a511f6156ff781cb60a1eb5cc59ca1 Mon Sep 17 00:00:00 2001
From: jakzar <jakzar2@st.amu.edu.pl>
Date: Sun, 26 May 2024 14:39:05 +0200
Subject: [PATCH] This: -fixed line where device for training was always cuda
 device

---
 neuralnetwork.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/neuralnetwork.py b/neuralnetwork.py
index 604accf..709d320 100644
--- a/neuralnetwork.py
+++ b/neuralnetwork.py
@@ -14,7 +14,7 @@ fertilizer = {labels[0]: 'kompost', labels[1]: 'saletra amonowa', labels[2]: 'su
 
 torch.manual_seed(42)
 
-device = torch.device('cuda')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def getTransformation():
     transform=transforms.Compose([