use reshape

visual.py

@@ -70,7 +70,7 @@ class Attention(nn.Module):
         out = attention_fn_default(
             q, k, v
         )
-        output = self.dense(out.transpose(1, 2).view(B, L, -1))
+        output = self.dense(out.transpose(1, 2).reshape(B, L, -1))
         output = self.output_dropout(output)
         return output