pytorch_discuss_mhr_issue

import torch
import torch.nn as nn
import torch.optim as optim

from torch.utils.data import Dataset, DataLoader


device = 'cpu'

class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.randn(100, 5)
        self.target = torch.empty(100, dtype=torch.long).random_(2)
        
    def __getitem__(self, index):
        x = self.data[index]
        y = self.target[index]
        
        return x, y

    def __len__(self):
        return len(self.data)


model = nn.Sequential(
    nn.Linear(5, 5),
    nn.ReLU(),
    nn.Linear(5, 2),
    nn.LogSoftmax(dim=1)
)
model.to(device)

criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-1)

dataset = MyDataset()
loader = DataLoader(dataset, batch_size=100, num_workers=2, shuffle=True)

epochs = 100
for epoch in range(epochs):
    for batch_idx, (data, target) in enumerate(loader):
        optimizer.zero_grad()
        data = data.to(device)
        target = target.to(device)
        
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
    print('Epoch {}, loss {}'.format(epoch, loss.item()))
    
    
# ===========================================================
    
    
def step(batch, model, criterion, optimizer=None):
    # let go of old gradients
    model.zero_grad()
    
    x = batch[0].to(device)
    y = batch[1].to(device)

    ## Forward Pass ##
    predictions = model(x)

    ## Calculate Loss ##
    loss = criterion(predictions, y)

    if optimizer is not None:
        # backward pass + optimize
        loss.backward()
        optimizer.step()
        
    return loss


def train_model(model=None, lr=0.01):
    criterion = nn.NLLLoss()
    params = list(filter(lambda p: p.requires_grad, model.parameters()))
    optimizer = torch.optim.Adam(params=params, lr=lr)

    for epoch in range(1, epochs+1):
        for i, batch in enumerate(loader):
            loss = step(batch, model, criterion, optimizer=optimizer)
        print('Epoch {}, loss {}'.format(epoch, loss.item()))


model = nn.Sequential(
    nn.Linear(5, 5),
    nn.ReLU(),
    nn.Linear(5, 2),
    nn.LogSoftmax(dim=1)
)
model.to(device)

train_model(model, lr=1e-1)