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

Clay-foundation / model

The Clay Foundation Model (in development)

https://clay-foundation.github.io/model/

Apache License 2.0

299 stars 38 forks source link

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

Closed weiji14 closed 8 months ago

weiji14 commented 8 months ago

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

[x] 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!
[x] 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 commented 8 months ago

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.