From 4db7338d11d890960ad0181416577dd916c7e730 Mon Sep 17 00:00:00 2001 From: Marius Brehler Date: Thu, 12 Dec 2024 12:22:54 +0000 Subject: [PATCH] [sharktank] Use reshape instead of view Fixes `RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.` by using `reshape`. which occurs when running `tests/layers/mmdit_test.py` with Python 3.11.10 and `torch==2.5.1+cpu`. --- sharktank/sharktank/layers/mmdit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sharktank/sharktank/layers/mmdit.py b/sharktank/sharktank/layers/mmdit.py index 0c970ab35..1557883ae 100644 --- a/sharktank/sharktank/layers/mmdit.py +++ b/sharktank/sharktank/layers/mmdit.py @@ -41,7 +41,7 @@ def attention(q, k, v, pe): q=q, k=k, v=v, a=None, is_causal=True, scale=None ) x = ops.permute(x, (0, 2, 1, 3)) - x = x.view(x.shape[0], x.shape[1], -1) + x = x.reshape(x.shape[0], x.shape[1], -1) return x