[Bug] When I tried YOLOXPOSE, the GPU graphics memory was insufficient.

Prerequisite

[X] I have searched Issues and Discussions but cannot get the expected help.
[X] The bug has not been fixed in the latest version(https://github.com/open-mmlab/mmpose).

Environment

Package Version Editable project location

addict 2.4.0 aliyun-python-sdk-core 2.15.1 aliyun-python-sdk-kms 2.16.2 attrs 23.2.0 certifi 2024.2.2 cffi 1.16.0 charset-normalizer 3.3.2 chumpy 0.70 click 8.1.7 colorama 0.4.6 contourpy 1.2.1 coverage 7.5.0 crcmod 1.7 cryptography 42.0.5 cycler 0.12.1 Cython 3.0.10 exceptiongroup 1.2.1 filelock 3.13.4 flake8 7.0.0 fonttools 4.51.0 fsspec 2024.3.1 idna 3.7 importlib_metadata 7.1.0 importlib_resources 6.4.0 iniconfig 2.0.0 intel-openmp 2021.4.0 interrogate 1.7.0 isort 4.3.21 Jinja2 3.1.3 jmespath 0.10.0 json-tricks 3.17.3 kiwisolver 1.4.5 Markdown 3.6 markdown-it-py 3.0.0 MarkupSafe 2.1.5 matplotlib 3.8.4 mccabe 0.7.0 mdurl 0.1.2 mkl 2021.4.0 mmcv 2.1.0 mmdet 3.2.0 mmengine 0.10.4 mmpose 1.3.1 d:\xuxinzhi\mmpose-main model-index 0.1.11 mpmath 1.3.0 munkres 1.1.4 networkx 3.2.1 numpy 1.26.4 opencv-python 4.9.0.80 opendatalab 0.0.10 openmim 0.3.9 openxlab 0.0.38 ordered-set 4.1.0 oss2 2.17.0 packaging 24.0 pandas 2.2.2 parameterized 0.9.0 pillow 10.3.0 pip 23.2.1 platformdirs 4.2.1 pluggy 1.5.0 py 1.11.0 pycocotools 2.0.7 pycodestyle 2.11.1 pycparser 2.22 pycryptodome 3.20.0 pyflakes 3.2.0 Pygments 2.17.2 pyparsing 3.1.2 pytest 8.1.2 pytest-runner 6.0.1 python-dateutil 2.9.0.post0 torch 2.2.0+cu121 torchaudio 2.2.0+cu121 torchvision 0.17.0+cu121 tqdm 4.65.2 typing_extensions 4.11.0 tzdata 2024.1 urllib3 1.26.18 wheel 0.41.2 xdoctest 1.1.3 xtcocotools 1.14.3 yapf 0.40.2 zipp 3.18.1

Reproduces the problem - code sample

class Swish(nn.Module): """Swish Module.

This module applies the swish function:

.. math::
    Swish(x) = x * Sigmoid(x)

Returns:
    Tensor: The output tensor.
"""

def __init__(self):
    super().__init__()

def forward(self, x: torch.Tensor) -> torch.Tensor:
    return x * torch.sigmoid(x)

Reproduces the problem - command or script

python tools/train.py configs/body_2d_keypoint/yoloxpose/coco/yoloxpose_s_8xb32-300e_coco-640.py

Reproduces the problem - error message

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 134.00 MiB. GPU 0 has a total capacity of 24.00 GiB of which 0 bytes is free. Of the allocated memory 21.09 GiB is allocated by PyTorch, and 2.02 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Additional information

No response

open-mmlab / mmpose