This: -fixed line where device for training was always cuda device

2024-05-26 14:39:05 +02:00 · 2024-05-26 14:39:05 +02:00 · 34bfe78218
parent 26c99ddc55
commit 34bfe78218
1 changed files with 1 additions and 1 deletions
--- 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([