incompatible type of tolerance in EMRegistration when gpu is used

[mireczech]

First of all, I would like to thank you for an awesome job with the library! I find a lot of use for it these days, so very much appreciated. I found one bug that keeps coming up, which should have an easy fix. When self.sigma2 is smaller than zero, it should be set to a value relative to the self.tolerance (see here). However, when GPU is used the type of self.sigma2 is changed by this operation to plain float and the script fails (see below). I would suggest to set the self.tolerance when initializing to be of the same type as self.sigma2 (see here). I can also create a pull request for this myself, if you prefer:)

Wish you a nice day!

The error:

  File "/home/miro/anaconda3/envs/SSM_anatomically_constrained/lib/python3.9/site-packages/torchcpd/emregistration.py", line 211, in register
    self.iterate()
  File "/home/miro/anaconda3/envs/SSM_anatomically_constrained/lib/python3.9/site-packages/torchcpd/emregistration.py", line 252, in iterate
    self.maximization()
  File "/home/miro/anaconda3/envs/SSM_anatomically_constrained/lib/python3.9/site-packages/torchcpd/emregistration.py", line 261, in maximization
    self.update_variance()
  File "/home/miro/anaconda3/envs/SSM_anatomically_constrained/lib/python3.9/site-packages/torchcpd/rigid_registration.py", line 135, in update_variance
    (2 * self.sigma2) + self.D * self.Np / 2 * torch.log(self.sigma2)
TypeError: log(): argument 'input' (position 1) must be Tensor, not float
make: *** [Makefile:61: data/registered] Error 1