How to run cosmos 1.0 7b text2world model on six 6000 Ada GPU Cards (each of which has 48GB memory)

Hi,

We’re trying to run cosmos 1.0 7b text2world model on six 6000 Ada GPU Cards (each of which has 48GB memory), but the results show
"
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 128.00 MiB. GPU 0 has a total capacity of 47.51 GiB of which 8.81 MiB is free. Process 8224 has 47.48 GiB memory in use. Of the allocated memory 46.99 GiB is allocated by PyTorch, and 12.60 MiB 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.
"

We used the following example to run our case.

PROMPT=“A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves.
The robot’s metallic body gleams under the bright, even lighting, highlighting its futuristic design and intricate joints.
A glowing blue light emanates from its chest, adding a touch of advanced technology. The background is dominated by rows of boxes,
suggesting a highly organized storage system. The floor is lined with wooden pallets, enhancing the industrial setting.
The camera remains static, capturing the robot’s poised stance amidst the orderly environment, with a shallow depth of
field that keeps the focus on the robot while subtly blurring the background for a cinematic effect.”

Example using 7B model

PYTHONPATH=$(pwd) python cosmos1/models/diffusion/inference/text2world.py
–checkpoint_dir checkpoints
–diffusion_transformer_dir Cosmos-1.0-Diffusion-7B-Text2World
–prompt “$PROMPT”
–offload_prompt_upsampler
–video_save_name Cosmos-1.0-Diffusion-7B-Text2World

We’re sure we have enough gpu memory to run this model,
Any comments to fix this problem are highly appreciated.