Update modeling_super_linear.py

modeling_super_linear.py

@@ -476,13 +476,14 @@ class superLinear(nn.Module):
                 ar_x = torch.cat([ar_x, ar_out], dim=1)[:, -self.seq_len:]
             out = torch.cat(outputs, dim=1)[:, :self.inf_pred_len]
             
-        print(F"out1 :{out.shape}")
         if len(x_enc.shape) > 2:
-            result =  out
-        else:
             out = out.reshape(B, V, out.shape[-1])
             print(F"out2 :{out.shape}")
             result = out.permute(0, 2, 1)
+        else:
+            print(F"out1 :{out.shape}")
+            result =  out
+
         
         if get_prob:
             expert_probs = expert_probs.reshape(B, V, expert_probs.shape[-1])