tianyi
/
auditory-visual
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

choose device automatically

This commit is contained in:
ClF3 2024-10-16 10:24:43 +08:00
parent 6d7f5d5c23
commit bbeb915b02
2 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored Normal file
View File

@ -0,0 +1,3 @@
*
!*.py
!.gitignore

2
CIFAR10_playground.py
View File

@ -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()