Simplification

high_order_layers_torch/Basis.py

@@ -259,7 +259,7 @@ def __call__(self, x):
 
         # These are the outputs, but they need to be in a sparse tensor
         # so they work with everything, do dense for now.
-        out = self._expand(x_in)
+        out = self._expand(x_in).permute(1, 2, 3, 4, 0)
 
         mat = torch.zeros(
             x.shape[0],
@@ -279,8 +279,6 @@ def __call__(self, x):
         # This needs to be
         windex = torch.div(torch.arange(wrange.numel()), self._n, rounding_mode="floor")
 
-        out = out.permute(1, 2, 3, 4, 0)
-
         mat_trans = mat.reshape(-1, self._variables)
         mat_trans[windex, wrange.view(-1)] = out.flatten()
         mat = mat_trans.reshape(