def METHOD_NAME():
    input_ids = torch.randint(tokenizer.vocab_size, (256,), device=torch.cuda.current_device())
    attn_mask = torch.ones_like(input_ids)
    return {'input_ids': input_ids, 'attention_mask': attn_mask}