How can I assign GPU allocation?

[nivibilla]

Im trying to shard the model into 4 parts over my 4 gpus but I can't seem to do it in-line.

config = ExLlamaConfig(model_config_path)          
config.model_path = model_path                         
config.max_seq_len = 1024 + 15                        
config.auto_map = [4,4,4,4]                             # This line doesn't work

model = ExLlama(config)                        
tokenizer = ExLlamaTokenizer(tokenizer_path)            

BATCH_SIZE = 16

cache = ExLlamaCache(model, batch_size=BATCH_SIZE)     
generator = ExLlamaGenerator(model, tokenizer, cache)  

[nivibilla]

Actually ignore this, it works, I just didn't have a large enough batch size