I am trying to do a random elastic deformation on the GPU during the training step within a Pytorch Lightning distributed environment. On creation of the Rand3DElasticd instance I am leaving the device parameter as default (None), because at this point I am not aware of the device on which the training will be executed. For Rand3DElastic there is already a lookup of the device of the supplied tensor in __call__, however, in Rand3DElasticd the configured device of Rand3DElastic will be used, which is undefined. Thus, the grid instantiation and resampling will always happen on the CPU.