Move criterion to selected device

As the criterion is not part of the model but instead a hyperparameter
due to the use of configuration files, it is not moved to the GPU if
selected as a device. We therefore manually move the criterion to the proper
device, which is bad practice when using lightning but works.