I have a pytorch lightning model on tabular data, a mix of categorical and continuous.
On Google Colab, switching from CPU to GPU only increases speed by 2x. (batch_size=1024 so it’s not small batches)
I suspect this is because I am doing a lot of matrix slicing, which is reducing locality. Here is a snippet from my forward method:
    def forward(self, x):
        xcat = []
        for i, (start, end) in enumerate(self.cat_ranges):
            xcat.append(self.embeddings[i](x[:,start:end]))
        xcat = torch.cat(xcat, 1)
        xcat = self.emb_dropout_layer(xcat)

        xcont = x[:,self.cont_range[0]:self.cont_range[1]]

Is this sort of slicing the reason I can’t speed up a lot on GPU? If so, how do I resolve it? Pre-slice the matrices so they are not views?