Training time is 14x in comparison to Keras

aksve: <p>I am tuning a model with multiple sequences to do a binary classification. Below is my model.</p>
<p>class LSTM_model(nn.Module):<br>
def <strong>init</strong>(self,vocab_size,hidden_size,add_feats,embed_size):<br>
super(LSTM_model, self).<strong>init</strong>()<br>
<span class="hashtag">#Model</span> Layers<br>
self.embedding0 = nn.Embedding(vocab_size+1,embed_size,padding_idx=def_val)<br>
self.embedding1 = nn.Embedding(vocab_size+1,embed_size,padding_idx=def_val)<br>
self.embedding2 = nn.Embedding(vocab_size+1,embed_size,padding_idx=def_val)<br>
self.embedding3 = nn.Embedding(vocab_size+1,embed_size,padding_idx=def_val)<br>
self.lstm0 = nn.LSTM(embed_size+add_feats, hidden_size,batch_first=True)#+1 is done to accomodate additional feature i.e. ts<br>
self.lstm1 = nn.LSTM(embed_size+add_feats, hidden_size,batch_first=True)<br>
self.lstm2 = nn.LSTM(embed_size+add_feats, hidden_size,batch_first=True)<br>
self.lstm3 = nn.LSTM(embed_size+add_feats, hidden_size,batch_first=True)<br>
self.fc1 = nn.Linear(hidden_size*4,64)<br>
self.fc2 = nn.Linear(64,32)<br>
self.fc3 = nn.Linear(32,1)<br>
self.drop = nn.Dropout(0.3)<br>
<span class="hashtag">#Activations</span><br>
self.relu = nn.ReLU()<br>
self.sigmoid = nn.Sigmoid()</p>
<pre><code>def forward(self,a,b,c,d,ts0,ts1,ts2,ts3):
    #Initialize the Embedding Layer
    x0 = self.embedding0(a.long())
    x1 = self.embedding1(b.long())
    x2 = self.embedding2(c.long())
    x3 = self.embedding3(d.long())
    #Numerical Features for Processing data 
    co0 = torch.cat([x0,ts0.unsqueeze(2)],2)
    co1 = torch.cat([x1,ts1.unsqueeze(2)],2)
    co2 = torch.cat([x2,ts2.unsqueeze(2)],2)
    co3 = torch.cat([x3,ts3.unsqueeze(2)],2)
    #Run 4 separate LSTM's
    x0 , (h0,c0) = self.lstm0(co0)
    x1 , (h1,c1) = self.lstm1(co1)
    x2 , (h2,c2) = self.lstm2(co2)
    x3 , (h3,c3) = self.lstm3(co3)
    #Transform the matrics
    h0 = h0[-1,:,:]
    h1 = h1[-1,:,:]
    h2 = h2[-1,:,:]
    h3 = h3[-1,:,:]
    #Transform the 4 layers
    f = torch.cat([h0,h1,h2,h3],1)
    x = self.relu(self.fc1(f))
    x = self.drop(x)
    x = self.relu(self.fc2(x))
    x = self.drop(x)
    x = self.sigmoid(self.fc3(x))
    
    return x    
</code></pre>
<p>This model is having the same number of parameters as compared to the keras model. Datasize is 100,000 rows and I am running on the GPU.</p>
<p>Dataloader code:</p>
<p>import torch.utils.data as data_utils<br>
dataset = torch.utils.data.TensorDataset(padded_events[events[0]],padded_events[events[1]],<br>
padded_events[events[2]],padded_events[events[3]],padded_ts[ts[0]],padded_ts[ts[1]],<br>
padded_ts[ts[2]],padded_ts[ts[3]],Y)</p>
<p>loader = torch.utils.data.DataLoader(<br>
dataset,<br>
batch_size=512,<br>
num_workers=0)</p>
<p>Model Training:</p>
<p><span class="hashtag">#Model</span><br>
model = LSTM_model(vocab_size=def_val,hidden_size=32,add_feats=1,embed_size=50)<br>
<span class="hashtag">#Loss</span><br>
loss_fn = nn.BCELoss()<br>
<span class="hashtag">#Optimizer</span><br>
optimizer = torch.optim.Adam(model.parameters(),lr=0.01)</p>
<p>%%time<br>
n_epochs = 1 # or whatever<br>
losses = []</p>
<p>for epoch in range(n_epochs):</p>
<pre><code>permutation = torch.randperm(df.shape[0])

for batch_idx, (a,b,c,d,ts0,ts1,ts2,ts3,y) in enumerate(loader):
    
    y = y.view(-1,1).float()
    a = a.clone().detach().requires_grad_(False)

    # in case you wanted a semi-full example
    outputs = model.forward(a,b,c,d,ts0,ts1,ts2,ts3)
    loss = loss_fn(outputs,y)
    print (outputs.requires_grad,y.requires_grad)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
  
    
    if batch_idx%10 == 0:
        losses.append(loss)
        log = open("/home/aksve/cs_experiment/data/write_output.txt", "a")
        log.write("epoch {}.\tbatch {}.\tloss : {}".format(epoch,batch_idx,loss))
        log.flush()
        log.close()
    
    
    if batch_idx%10 == 0:
        losses.append(loss)
        print("epoch {}.\tbatch {}.\tloss : {}".format(epoch,batch_idx,loss))
</code></pre>
<p>one epoch in keras takes 58 seconds whereas it taken 14 min in torch. Is there anything inherently wrong here. If I increase the batch size to 8192, torch time reduces to 68 seconds. Unable to understand why is it happening</p>