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

from how do I eliminate the for-loop in this pytorch code