I made this training code part, implementing my previous mnist code.
It work good and fast(it takes about 12sec per 1 epoch)
In that case even though the printing part was out side of loading loop and used cpu.
MNIST has 60000 training image data so I don’t think this code(batch size is 100) below is getting the data 60000*100 times. or is it???
# dataset loader
data_loader = torch.utils.data.DataLoader(dataset=mnist_train,
                                          batch_size=batch_size,   #batch_size = 100
                                          shuffle=True,
                                          drop_last=True)

for epoch in range(training_epochs):
    avg_cost = 0
    total_batch = len(data_loader)

    for X, Y in data_loader:
        # reshape input image into [batch_size by 784]
        # label is not one-hot encoded
        X = X.view(-1, 28 * 28).to(device)
        Y = Y.to(device)

        optimizer.zero_grad()
        hypothesis = linear(X)
        cost = criterion(hypothesis, Y)
        cost.backward()
        optimizer.step()

        avg_cost += cost / total_batch

    print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.9f}'.format(avg_cost))

print('Learning finished')