Cuda out of memory while using Trainer API

Could this be related to this issue?