[PyTorch] Inference and fine-tunning with batch size different of 1

[rscarvalho90]

Hi!

Thanks for the "translation" of mUSE to PyTorch.

Did you test the PyTorch model while fine tuning?

In my case, I'm trying to fine tune the model adding it as a layer in a custom model.
However, apparently the MUSE class is not prepared to receive batches of inputs.

This is the error (100 is the batch size):

architecture.py:26, in PositionalEncoding.forward(self, x)
     25 def forward(self, x: torch.Tensor) -> torch.Tensor:
---> 26     x = x + self.pe[: x.size(0)]
     27     return x

RuntimeError: The size of tensor a (512) must match the size of tensor b (100) at non-singleton dimension 1

The same occurs when I'm trying to fine tune the original model directly, without add it as a layer in a custom model.

The model is being fed by a Tensor with shape [batch_size, 512].