Update modeling_super_linear.py

modeling_super_linear.py

@@ -475,10 +475,14 @@ class superLinear(nn.Module):
                 outputs.append(ar_out)
                 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}")
-        out = out.reshape(B, V, out.shape[-1])
-        print(F"out2 :{out.shape}")
-        result = out.permute(0, 2, 1)
+        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)
         
         if get_prob:
             expert_probs = expert_probs.reshape(B, V, expert_probs.shape[-1])