diff --git a/beginner_source/basics/buildmodel_tutorial.py b/beginner_source/basics/buildmodel_tutorial.py index 5cbaf60602..dd4f4560a6 100644 --- a/beginner_source/basics/buildmodel_tutorial.py +++ b/beginner_source/basics/buildmodel_tutorial.py @@ -75,7 +75,7 @@ def forward(self, x): # along with some `background operations `_. # Do not call ``model.forward()`` directly! # -# Calling the model on the input returns a 2-dimensional tensor with dim=0 corresponding to each output of 10 raw predicted values for each class, and dim=1 corresponding to the individual values of each output. +# Calling the model on the input returns a 2-dimensional tensor of shape ``(batch_size, 10)``, with dim=0 corresponding to each sample in the batch and dim=1 corresponding to the raw predicted values (logits) for each of the 10 classes. # We get the prediction probabilities by passing it through an instance of the ``nn.Softmax`` module. X = torch.rand(1, 28, 28, device=device)