MEMORY EXHAUSTED when running the ‘image observation example’ program on an NVIDIA GeForce RTX 4070 Ti

[WenzheHe]

I ran into a problem with insufficient memory when running the image observation example(serl/examples/async_drq_sim/async_drq_sim.py) on my NVIDIA GeForce RTX 4070 Ti.The error is as follows:

E1221 16:46:49.290140 30613 pjrt_stream_executor_client.cc:3085] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 1208256760 bytes.
jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 1208256760 bytes.

The error occurs when running code 'agent, update_info = agent.update_high_utd(batch, utd_ratio=1)'.

I checked my computer's memory usage, it's not running any other processes, and there's about 8G of free memory. I don't know how to fix this, any suggestions would be greatly appreciated, thanks!

[yoimika]

I met the same error on my 3070 8GB, but solved by reducing the batch size to 64. The original batch size is 256 which will cause the same error on 4090 24GB.