How to write a fast PyTorch training loop?
I've encountered a bad regression in performance after porting a machine learning algorithm from Keras to PyTorch.
The following construction in Keras:
model.compile(loss=&quot;binary_crossentropy&quot;, optimizer=SGD(learning_rate=0.01)) 
.
.
.
model.fit(states, actions)

is ~15x faster than the following verbose version in Pytorch:
def train_network(model, optimizer, train_loader,
                  num_epochs=16, pbar_update_interval=200, print_logs=True):

    criterion = nn.BCELoss()
    model.train()

    for i in range(num_epochs):
        for k, batch_data in enumerate(train_loader):
            optimizer.zero_grad()

            batch_x = batch_data[:, :-1]
            batch_y = batch_data[:, -1]

            y_pred = model(batch_x)
            loss = criterion(y_pred, batch_y.unsqueeze(1))
            loss.backward() 
            optimizer.step() 

optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)
train_data = torch.from_numpy(np.column_stack((states, actions)))
train_data = train_data.to(torch.float)
train_loader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=32)
train_network(model, optimizer, train_loader)

Am I writing the PyTorch code wrong?
The variable model in the PyTorch code is an nn.Module, and states and actions are NumPy arrays. The code is running on 4 CPU cores.

It seems that keras.model.fit has default parameter value epochs=1. So this exactly explains the performance factor.
