Skip to content

Commit 9e8a306

Browse files
Update CIFAR10 tutorial to use torchvision v2 transforms and torch.ac… (#3882)
…celerator Fixes #3878 ## Description Updated CIFAR10 tutorial APIs and device handling - Replace deprecated ToTensor with v2 transforms pipeline - Use transforms.v2 namespace consistently - Replace CUDA-only device detection with torch.accelerator - Switch model checkpoint extension from .pth to .pt - Update GitHub links from master to main <img width="1815" height="615" alt="image" src="https://github.com/user-attachments/assets/1910355b-139a-48e0-935f-f3392c5ae140" /> ## Checklist <!--- Make sure to add `x` to all items in the following checklist: --> - [ ] The issue that is being fixed is referred in the description (see above "Fixes #ISSUE_NUMBER") - [ ] Only one issue is addressed in this pull request - [ ] Labels from the issue that this PR is fixing are added to this pull request - [ ] No unnecessary issues are included into this pull request. Co-authored-by: sekyondaMeta <127536312+sekyondaMeta@users.noreply.github.com>
1 parent cb04c52 commit 9e8a306

1 file changed

Lines changed: 10 additions & 9 deletions

File tree

beginner_source/blitz/cifar10_tutorial.py

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@
5757
"""
5858
import torch
5959
import torchvision
60-
import torchvision.transforms as transforms
60+
from torchvision.transforms import v2
6161

6262
########################################################################
6363
# The output of torchvision datasets are PILImage images of range [0, 1].
@@ -69,9 +69,10 @@
6969
# BrokenPipeError or RuntimeError related to multiprocessing, try setting
7070
# the num_worker of torch.utils.data.DataLoader() to 0.
7171

72-
transform = transforms.Compose(
73-
[transforms.ToTensor(),
74-
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
72+
transform = v2.Compose([
73+
v2.ToImage(),
74+
v2.ToDtype(torch.float32, scale=True),
75+
v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
7576

7677
batch_size = 4
7778

@@ -191,7 +192,7 @@ def forward(self, x):
191192
########################################################################
192193
# Let's quickly save our trained model:
193194

194-
PATH = './cifar_net.pth'
195+
PATH = './cifar_net.pt'
195196
torch.save(net.state_dict(), PATH)
196197

197198
########################################################################
@@ -302,7 +303,7 @@ def forward(self, x):
302303
# Let's first define our device as the first visible cuda device if we have
303304
# CUDA available:
304305

305-
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
306+
device = torch.device(torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else 'cpu')
306307

307308
# Assuming that we are on a CUDA machine, this should print a CUDA device:
308309

@@ -355,9 +356,9 @@ def forward(self, x):
355356
# - `Discuss PyTorch on the Forums`_
356357
# - `Chat with other users on Slack`_
357358
#
358-
# .. _Train a state-of-the-art ResNet network on imagenet: https://github.com/pytorch/examples/tree/master/imagenet
359-
# .. _Train a face generator using Generative Adversarial Networks: https://github.com/pytorch/examples/tree/master/dcgan
360-
# .. _Train a word-level language model using Recurrent LSTM networks: https://github.com/pytorch/examples/tree/master/word_language_model
359+
# .. _Train a state-of-the-art ResNet network on imagenet: https://github.com/pytorch/examples/tree/main/imagenet
360+
# .. _Train a face generator using Generative Adversarial Networks: https://github.com/pytorch/examples/tree/main/dcgan
361+
# .. _Train a word-level language model using Recurrent LSTM networks: https://github.com/pytorch/examples/tree/main/word_language_model
361362
# .. _More examples: https://github.com/pytorch/examples
362363
# .. _More tutorials: https://github.com/pytorch/tutorials
363364
# .. _Discuss PyTorch on the Forums: https://discuss.pytorch.org/

0 commit comments

Comments
 (0)