How to use a designated GPU for training?

[shishijier]

I have 8 GPUs and I want to use one of them for training, but I haven't found the relevant settings.

from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Trainer

if __name__=='__main__':

    with Run().context(RunConfig(nranks=1, experiment="msmarco")):

        config = ColBERTConfig(
            bsize=16,
            root="/path/to/experiments",
        )
        trainer = Trainer(
            triples="triples.tsv",
            queries="queries.tsv",
            collection="collectios.tsv",
            config=config,
        )

        checkpoint_path = trainer.train()

        print(f"Saved checkpoint to {checkpoint_path}...")

[Hannibal046]

maybe setting CUDA_VISIBLE_DEVICES=x?

[adrienB134]

I think that if you set nranks to 8 you should be able to use your 8 GPUS.