TheseusLayer.to improvements: Return itself and support cuda

[tvercaut]

🚀 Feature

It would be great to be able to call

theseus_optim = th.TheseusLayer(optimizer).to(torch.device('cuda'))

as well as

theseus_optim = th.TheseusLayer(optimizer, device=torch.device('cuda'))

Currently this fails for two reasons:

• TheseusLayer.to returns None and this the one-liner needs to be split in 2: theseus_optim = th.TheseusLayer(optimizer); theseus_optim.to(torch.device('cuda'))
• Using 'cuda' fails as this is mapped to cuda:0 in PyTorch tensors and the comparison then fails in Objective.update
  

  theseus/theseus/core/objective.py

  Lines 775 to 780 in 9a117fd

  	if tensor.device != self.device or tensor.dtype != self.dtype:
  	raise ValueError(
  	f"Attempted to update variable {var_name} with a "
  	f"({tensor.device},{tensor.dtype}) tensor, which is inconsistent "
  	f"with objective's expected ({self.device},{self.dtype})."
  	)

Motivation

This would make the use of Theseus more convenient.

Pitch

See above

Alternatives

See above

Additional context

theseus/theseus/theseus_layer.py

Lines 137 to 140 in 9a117fd

	# Applies to() with given args to all tensors in the objective
	def to(self, *args, **kwargs):
	super().to(*args, **kwargs)
	self.objective.to(*args, **kwargs)

Other issues with to have been discussed in:

• Add option to pass device when constructing a new Objective #281
• Baspacho device doesn't accept "cuda:0" #597

[luisenp]

This is a good suggestion, thanks for bringing it up! One question is about the expected behavior when passing torch.device("cuda) . Should we automatically pass to "cuda:0" on our side? As you say, this seems to be the case for torch tensors, but for some reason torch.device() doesn't do this automatically. Not sure if this is an oversight on Pytorch's side or a deliberate choice.

[luisenp]

Added feature in #623.