Content deleted Content added
TomRitchford (talk | contribs) →Example: Simplify call to super().__init__(). |
TomRitchford (talk | contribs) →Example: Make comments PEP-8 |
||
Line 97:
from torch import nn # Import the nn sub-module from PyTorch
class NeuralNetwork(nn.Module): # Neural networks are defined as classes
def __init__(self): # Layers and variables are defined in the __init__ method
super().__init__() # Must be in every network.
self.flatten = nn.Flatten() #
self.linear_relu_stack = nn.Sequential( #
nn.Linear(28*28, 512), # Linear Layers have an input and output shape
nn.ReLU(), # ReLU is one of many activation functions provided by nn
nn.Linear(512, 512),
nn.ReLU(),
Line 109:
)
def forward(self, x): # This function defines the forward pass.
x = self.flatten(x)
logits = self.linear_relu_stack(x)
| |||