Pytorch for loop vs reshaping "Time Distributed" outputs same forward results, but trains differently

I have the torch model that processes sentences and does so using a for loop.

I recently rewrote the forward method to instead reshape and process everything all at once and then reshape the results again (similar to TimeDistributed in Tensorflow/Keras).

When I run each forward method, they output the same results. However, when I train each model using the same specifications, they seem to perform differently.

Does how things are reshaped affect how the gradients are distributed? Or does anyone have any insight on what is going on here?

Here is the code for comparison:

Original:

masks = self.mask_pad(contexts, 0).transpose(0,1)

if self.use_pos == True:
    x = self.pos_enc(self.pos_att(self.emb(contexts))).transpose(0,1)
else:
    x = self.emb(contexts).transpose(0,1)

res = []
for xi, mask in zip(x, masks):
    for layer in self.context_encoder:
        xi = layer(xi, mask)
    mask_value = mask.squeeze().unsqueeze(-1).float()

    div_val = torch.sum(mask_value, dim=1)
    div_val[div_val==0] = 1 #replace div_vals of 0 with 1, to avoid divide by 0 error.  these should be 0 vecs anyway
   
    res += [torch.sum(xi * mask_value, dim=1) / div_val]

res = torch.stack(res).transpose(0,1)

Faster:

masks = self.mask_pad(contexts, 0).transpose(0,1)

if self.use_pos == True:
    x = self.pos_enc(self.pos_att(self.emb(contexts))).transpose(0,1)
else:
    x = self.emb(contexts).transpose(0,1)

        
#print('x '+str(x.size()))  
#print('masks '+str(masks.size()))  

#print('--')

x_reshape = x.contiguous().view(-1, x.size(-2), x.size(-1))  # (samples * timesteps, input_size)
mask_reshape = masks.contiguous().view(-1, masks.size(-3), masks.size(-2), masks.size(-1))  # (samples * timesteps, input_size)

for layer in self.context_encoder:
    x_reshape = layer(x_reshape, mask_reshape)

mask_value = mask_reshape.squeeze().unsqueeze(-1).float()

div_val = torch.sum(mask_value, dim=1)
div_val[div_val==0] = 1 #replace div_vals of 0 with 1, to avoid divide by 0 error.  these should be 0 vecs anyway
res = torch.sum(x_reshape * mask_value, dim=1) / div_val

print(res.size()) #torch.Size([4160, 400])
res = res.view(-1, x.size(1), res.size(-1))  # (timesteps, samples, output_size)
res = res.transpose(0,1)