title: how do I eliminate the for-loop in this pytorch code
body: <p>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.</p>
  <pre><code>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
  </code></pre>