Refactor model for multi-device usage and easier disabling of masking

[weiji14]

A couple of small changes to model_clay.py after initial implementation at #47.

• Handle layers across cpu/gpu devices a bit better, to prevent errors like RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
• Let the masking in the Masked Autoencoder be dynamic, so that setting self.model.encoder.mask_ratio = 0 would disable the masking.

First one is cherry-picked from @srmsoumya's ddp branch. Second one will be helpful for the embedding generation later on.

[weiji14]

@srmsoumya, I feel like we can also remove the .to(device) call here?

Suggested change

	# Move position & band encoding to the device
	self.pos_encoding = self.pos_encoding.to(patches.device)
	self.band_encoding = self.band_encoding.to(patches.device)

[weiji14]

Ah, I think I found the core issue. The pos_encoding and band_encoding tensors are declared in the __init__ method, before Lightning has managed to cast them to the correct device.

model/src/model_clay.py

Lines 89 to 95 in ee74c91

	# Fix the position & band embedding to sine & cosine functions
	self.pos_encoding = posemb_sincos_2d(
	h=image_size // patch_size, w=image_size // patch_size, dim=pos_dim
	) # [L D/2]
	self.band_encoding = posemb_sincos_1d(
	length=self.num_group_patches, dim=band_dim
	) # [G D/2]

Ideally, we'll need to define the device at this point. Let me see if there's a way to do that.

Edit: Yep, found the instructions - https://lightning.ai/docs/pytorch/2.1.0/starter/converting.html#remove-any-cuda-or-to-device-calls

[weiji14]

Ok, went with using self.register_buffers at commit 25f83e6. Have set persistent=False since those pos_encoding and band_encoding tensors are fixed, and won't need to be saved to the state_dict.

[weiji14]

Ok, gonna merge this patch in directly since others may benefit from the improve cuda/cpu device handling. Plus I'll need these patches for the embedding factory script.