How to use Clay foundation model on GPU?

[sri-ram-swaminathan]

Hi all,

First of all, thank you for the awesome model :)

I'm trying to do bio-diversity classification on RGB-NIR imagery and tested out the Clay model on my CPU -- which worked without error.

I now want to run the script on my training data and I'm using a GPU for this. Whenever I try to run it like this, I get the following error:

call._call_and_handle_interrupt(
File "/path/to/venv/lib/python3.12/site-packages/lightning/pytorch/trainer/call.py", line 49, in _call_and
self._run(model, ckpt_path=ckpt_path, weights_only=weights_only)
File "/path/to/venv/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py", line 1121, in _run_sanity_check
return loop_run(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/path/to/venv/lib/python3.12/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 146, in _evaluation_step
File "/path/to/venv/lib/python3.12/site-packages/lightning/pytorch/trainer/call.py", line 329, in _call_st
^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^
File "/path/to/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in wrapped_call
File "/path/to/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
File "/path/to/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/path/to/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/path/to/project/src/myproject/models/models.py", line 67, in forward
File "/path/to/venv/lib/python3.12/site-packages/terratorch/models/backbones/clay_v1/modules.py", line 413, in forward
patches = self.add_encodings(
^^^^^^^^^^^^^^^^^^^
File "/path/to/venv/lib/python3.12/site-packages/terratorch/models/backbones/clay_v1/modules.py", line 379, in add_encodings
posemb_sincos_2d_with_gsd(
File "/path/to/venv/lib/python3.12/site-packages/terratorch/models/backbones/clay_v1/utils.py", line 23, in posemb_sincos_2d_with_gsd
omega = 1.0 / (temperature ** (2 * omega / dim)) *
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

The code I'm using is:

class Clay(nn.Module):
    def __init__(self, freeze_encoder: bool = True, num_classes: int = 3):
        super().__init__()
        self.encoder = BACKBONE_REGISTRY.build("timm_clay_v1_base", pretrained=True)
        slef.encoder = self.encoder.to("cuda") # Attempt to fix
        embd_dim = 768
        self.num_classes = num_classes
        if freeze_encoder:
            for param in self.encoder.parameters():
                param.requires_grad = False

        self.decoder = nn.Sequential(
            nn.Linear(embd_dim, embd_dim // 2),
            nn.ReLU(),
            nn.Linear(embd_dim // 2, self.num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.interpolate(x, size=(256, 256), mode="bicubic", align_corners=False) # Clay expects 256 x 256 images
        x = x[:, [2, 1, 0, 3], :, :] # (R,G,B,NIR) --> (B,G,R,NIR)
        datacube = {
            "pixels": x,
            "time":   torch.zeros(x.shape[0], 4, device=x.device),
            "latlon": torch.zeros(x.shape[0], 4, device=x.device),
            "gsd":    torch.tensor([10], device=x.device),
            "waves":  torch.tensor([0.485,0.560,0.660,0.835], device=x.device),
        }
        out = self.encoder.clay_encoder(datacube)
        return self.decoder(out[-1][:, 1:, :].mean(dim=1))

I've already tried forcing the encoder to the GPU but this doesn't help either

I know that the code works because I'm training other foundational models and they work fine without any changes

Could you please help me out?