def assign_gpu()

in mico/utils/utils.py [0:0]

• 8 lines of code
• 1 McCabe index (conditional complexity)

def assign_gpu(tokenizer_output, device='cuda'):
    """This function helps move the tokenizer output to GPU since we use the BERT model on GPU.
    """
    tokens_tensor = tokenizer_output['input_ids'].to(device)
    token_type_ids = tokenizer_output['token_type_ids'].to(device)
    attention_mask = tokenizer_output['attention_mask'].to(device)

    output = {'input_ids': tokens_tensor,
              'token_type_ids': token_type_ids,
              'attention_mask': attention_mask}
    return output