how do I eliminate the for-loop in this pytorch code
I have this custom pytorch module (below). It does exactly what I need; it just does it very slowly. What can I do to speed this up? I know that I'm not supposed to have a for-loop in there; it just wasn't clear how to do that divide operation without it. How do I broadcast the x tensor to the divide without that loop? I can move the back-weights to their own layer if that helps.
class StepLayer(nn.Module):
    def __init__(self):
        super(StepLayer, self).__init__()
        w = init_weights()
        self.front_weights = nn.Parameter(torch.DoubleTensor([w, w]).T, requires_grad=True)
        self.back_weights = nn.Parameter(torch.DoubleTensor([w]).T, requires_grad=True)
        

    def forward(self, x):
        # x shape is batch by feature
        results = []
        for batch in x:
            b = batch.divide(self.front_weights)
            b = torch.some_math_function(b)
            b = b.sum(dim=1)
            b = torch.some_other_math_function(b)
            b = b @ self.back_weights
            results.append(b)
        stack = torch.vstack(results)
        return stack


Below is a source code with shape after each step described (read code comments please).
I've assumed a few things like F=100, x=Bx2, front_weights=100x2, back_weights=100, you should be able to easily adjust it to your case.
class StepLayer(nn.Module):
    def __init__(self):
        super().__init__()
        F = 100
        # Notice I've added `1` dimension in front_weights
        self.front_weights = nn.Parameter(torch.randn(1, F, 2), requires_grad=True)
        self.back_weights = nn.Parameter(torch.randn(F), requires_grad=True)

    def forward(self, x):
        # x.shape == (B, 2)
        x = x.unsqueeze(dim=1)  # (B, 1, 2)
        x = x / self.front_weights  # (B, F, 2)
        # I just took some element-wise math function from PyTorch
        x = torch.sin(x)  # (B, F, 2)
        x = torch.sum(x, dim=-1)  # (B, F)
        x = torch.sin(x)  # (B, F)
        return x @ self.back_weights  # (B, )

        # results = []
        # for batch in x:
        #     # batch - (1, 2)
        #     b = batch.divide(self.front_weights)  # (F, 2)
        #     b = torch.some_math_function(b)  # (F, 2)
        #     b = b.sum(dim=1)  # (F, )
        #     b = torch.some_other_math_function(b)  # (F, )
        #     b = b @ self.back_weights  # (1, )
        #     results.append(b)
        # stack = torch.vstack(results)  # (B, )
        # return stack  # (B,)


layer = StepLayer()

print(layer(torch.randn(64, 2)).shape)

Main trick is to use 1 dimensions for broadcasting when necessary (division especially) and smart weight initialization so you don't have to do any transpose operations.
Other things

You might want to reconsider Double, float (as used above) is much faster, especially on CUDA and takes half the memory (and neural network should compensate for precision loss, if any).
Use half precision and mixed training if the speed is still an issue (float16 dtype instead of float32), but only on CUDA, see here to know more about automated mixed precision

