Replaces calls to .cuda with .to(torch_device) in tests (#25571)
* Replaces calls to `.cuda` with `.to(torch_device)` in tests `torch.Tensor.cuda()` is a pre-0.4 solution to changing a tensor's device. It is recommended to prefer `.to(...)` for greater flexibility and error handling. Furthermore, this makes it more consistent with other tests (that tend to use `.to(torch_device)`) and ensures the correct device backend is used (if `torch_device` is neither `cpu` or `cuda`). * addressing review comments * more formatting changes in Bloom test * `make style` * Update tests/models/bloom/test_modeling_bloom.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * fixes style failures --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
This commit is contained in:
@@ -497,13 +497,13 @@ class XGLMModelLanguageGenerationTest(unittest.TestCase):
|
||||
model_name = "facebook/xglm-564M"
|
||||
tokenizer = XGLMTokenizer.from_pretrained(model_name, use_fast=False, padding_side="left")
|
||||
|
||||
model = XGLMForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, use_cache=True).cuda()
|
||||
model = XGLMForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, use_cache=True).to(torch_device)
|
||||
model = model.eval()
|
||||
|
||||
batch = tokenizer(["Who are you?", "Joe Biden is the president of"], padding=True, return_tensors="pt")
|
||||
|
||||
input_ids = batch["input_ids"].cuda()
|
||||
attention_mask = batch["attention_mask"].cuda()
|
||||
input_ids = batch["input_ids"].to(torch_device)
|
||||
attention_mask = batch["attention_mask"].to(torch_device)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(input_ids, attention_mask=attention_mask)
|
||||
|
||||
Reference in New Issue
Block a user