Add a device map

[ideasbyjin]

Hey guys, good stuff on this -- I'd recommend allowing a device map feature so you can load ProtT5 weights directly to GPU.

This does have a dependency on accelerate though:

class ProtTransEmbedder(LightningModule):
    def __init__(
        self,
        weights_dir: Path,
        model_type: str,
        device_map: str = 'auto',
    ) -> None:
        ...
        if model_type == "bert":
            self.model = BertModel.from_pretrained(
                self.weights_dir, add_pooling_layer=False, device_map = device_map
            )
        elif model_type == "t5":
            self.model = T5EncoderModel.from_pretrained(self.weights_dir, device_map = device_map)

[henrykenlay]

Hi ideasbyjin, thank you for your interest and suggestion. My collaborator @fdreyer made a PR based on your snippet: https://github.com/Exscientia/abodybuilder3/pull/4. If this looks okay to you we can merge it :)