importos;os.environ['CUDA_VISIBLE_DEVICES']='2'# use 1 device onlyimporttorchfromtransformersimportOPTForCausalLMmodel_name='facebook/opt-350m'model=OPTForCausalLM.from_pretrained(model_name,device_map='cuda')# memory usage: 1688 MiBbatch_size=6seq_len=2048seq_ids=torch.zeros((batch_size,...