Your linear layer is not doing the same as a TimeDistributedDense in Keras. You are only using the last time step, and ditching everything else.
Have a look at my TimeDistributed wrapper here:

  
    
    
    Any PyTorch function can work as Keras' Timedistributed? vision
  
  
    Hey, 
I developed a PyTorch module that mimics the TimeDistributed wrapper of Keras a few days ago: 
import torch.nn as nn


class TimeDistributed(nn.Module):
    def __init__(self, module, batch_first=False):
        super(TimeDistributed, self).__init__()
        self.module = module
        self.batch_first = batch_first

    def forward(self, x):

        if len(x.size()) <= 2:
            return self.module(x)

        # Squash samples and timesteps into a single axis
        x_reshape = x.…