From bbeb915b021f82f653f6e19166a4e9f9255959a5 Mon Sep 17 00:00:00 2001
From: ClF3 <smh66ccff@outlook.com>
Date: Wed, 16 Oct 2024 10:24:43 +0800
Subject: [PATCH] choose device automatically

---
 .gitignore            | 3 +++
 CIFAR10_playground.py | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)
 create mode 100644 .gitignore

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..8a0a550
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+*
+!*.py
+!.gitignore
\ No newline at end of file
diff --git a/CIFAR10_playground.py b/CIFAR10_playground.py
index 7df1dc8..404637e 100755
--- a/CIFAR10_playground.py
+++ b/CIFAR10_playground.py
@@ -223,7 +223,7 @@ class BnDeepNet(nn.Module):
 model = DeepNet('tanh')
 # model = BnDeepNet('relu')
 
-device=torch.device("cuda")
+device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 criterion = nn.CrossEntropyLoss()