mobilevit_xs occured CUDA out of memory?

[HTCho1]

I test my code using pretrained mobilevit_xs and pretrained HRNet(Deep High-Resolution Representation Learning for Visual Recognition) by my dataset. I checked that mobilevit_xs have less x8 parameters than HRNet and my pretrained weight of mobilevit_xs (85MB) is smaller than pretrained weight of HRNet (254MB).

When i test the model using eval image, the model using mobilevit_xs occured CUDA out of memory on RTX 3070. However, HRNet model work normally.

Why do not work the mobilevit_xs model?
Eval image resolution is 1080 x 1920.

Is it because of key, query, value operations?

Error message is below.
Traceback (most recent call last):
File ".\streaming_points.py", line 213, in
writeVideo(args)
File ".\streaming_points.py", line 146, in writeVideo
model = build_model(args.weight_path)
File ".\streaming_points.py", line 114, in build_model
_ = model(dummy_tensor)
File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 224, in forward
x = self.mvit0
File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 161, in forward
x = self.transformer(x)
File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 89, in forward
x = attn(x) + x
File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 31, in forward
return self.fn(self.norm(x), **kwargs)
File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 70, in forward
dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
RuntimeError: CUDA out of memory. Tried to allocate 1.99 GiB (GPU 0; 8.00 GiB total capacity; 3.85 GiB already allocated; 1.53 GiB free; 3.87 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_s
plit_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF