# Any PyTorch function can work as Keras' Timedistributed?

**URL:** <https://discuss.pytorch.org/t/any-pytorch-function-can-work-as-keras-timedistributed/1346>\
**Category:** vision\
**Created:** [March 25, 2017, 10:35am UTC](https://discuss.pytorch.org/t/any-pytorch-function-can-work-as-keras-timedistributed/1346 "2017-03-25T10:35:07Z")\
**Posts on this page:** 1\
**Showing post:** 4

<div class="post-metadata">

**Author:** ![miguelvr](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/miguelvr/32/366_2.png) [@miguelvr](https://discuss.pytorch.org/u/miguelvr)\
**Post date:** [March 28, 2017, 10:24pm UTC](https://discuss.pytorch.org/t/any-pytorch-function-can-work-as-keras-timedistributed/1346/4 "2017-03-28T22:24:30Z")

</div>

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.contiguous().view(-1, x.size(-1)) # (samples * timesteps, input_size)

        y = self.module(x_reshape)

        # We have to reshape Y
        if self.batch_first:
            y = y.contiguous().view(x.size(0), -1, y.size(-1)) # (samples, timesteps, output_size)
        else:
            y = y.view(-1, x.size(1), y.size(-1)) # (timesteps, samples, output_size)

        return y
```

---

_[View the full topic](https://discuss.pytorch.org/t/any-pytorch-function-can-work-as-keras-timedistributed/1346)._