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([