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

[2823387967]

Prerequisite

• I have searched Issues and Discussions but cannot get the expected help.
• 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