CUDA Out of Memory when using depth_extractor in train_multi_bc_deter.py #26
Description
Traceback (most recent call last):
File "train_multi_bc_deter.py", line 404, in
trainer.train()
File "/home/wang/Desktop/visual_wholebody/high-level/learning/dagger_trainer.py", line 70, in train
self.single_agent_train()
File "/home/wang/Desktop/visual_wholebody/high-level/learning/dagger_trainer.py", line 153, in single_agent_train
self.agents.post_interaction(timestep=timestep, timesteps=self.timesteps)
File "/home/wang/Desktop/visual_wholebody/high-level/learning/dagger_rnn.py", line 265, in post_interaction
self._update(timestep, timesteps)
File "/home/wang/Desktop/visual_wholebody/high-level/learning/dagger_rnn.py", line 321, in _update
student_actions, _, _ = self.policy.act({"states": sampled_student_obs, **rnn_policy}, role="policy")
File "/home/wang/Desktop/visual_wholebody/third_party/skrl/skrl/models/torch/deterministic.py", line 88, in act
actions, outputs = self.compute(inputs, role)
File "train_multi_bc_deter.py", line 145, in compute
depth_feature = self.depth_extractor(images)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/home/wang/Desktop/visual_wholebody/high-level/modules/feature_extractor.py", line 128, in forward
images_compressed = self.image_compression(images)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/container.py", line 219, in forward
input = module(input)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 458, in forward
return self._conv_forward(input, self.weight, self.bias)
File "/home/wang/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 454, in _conv_forward
return F.conv2d(input, weight, bias, self.stride,
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.76 GiB. GPU 0 has a total capacity of 15.56 GiB of which 1.33 GiB is free. Including non-PyTorch memory, this process has 12.89 GiB memory in use. Of the allocated memory 5.77 GiB is allocated by PyTorch, and 155.23 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. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)