Then your Docker compose container should have:
image-name:
build:
context: .
devices:
- /dev/dri
- /dev/kfd
group_add:
- video
shm_size: 4G
environment:
- PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.9,max_split_size_mb:512
- PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.9,max_split_size_mb:512
- HSA_OVERRIDE_GFX_VERSION=11.0.0