wrap forward passes with torch.no_grad() (#19273)

This commit is contained in:
Partho
2022-10-04 19:43:22 +05:30
committed by GitHub
parent d6e920449e
commit a9782881a4

View File

@@ -627,6 +627,7 @@ class BigBirdModelIntegrationTest(unittest.TestCase):
model.to(torch_device) model.to(torch_device)
input_ids = torch.tensor([[20920, 232, 328, 1437] * 1024], dtype=torch.long, device=torch_device) input_ids = torch.tensor([[20920, 232, 328, 1437] * 1024], dtype=torch.long, device=torch_device)
with torch.no_grad():
outputs = model(input_ids) outputs = model(input_ids)
prediction_logits = outputs.prediction_logits prediction_logits = outputs.prediction_logits
seq_relationship_logits = outputs.seq_relationship_logits seq_relationship_logits = outputs.seq_relationship_logits
@@ -655,6 +656,7 @@ class BigBirdModelIntegrationTest(unittest.TestCase):
model.to(torch_device) model.to(torch_device)
input_ids = torch.tensor([[20920, 232, 328, 1437] * 512], dtype=torch.long, device=torch_device) input_ids = torch.tensor([[20920, 232, 328, 1437] * 512], dtype=torch.long, device=torch_device)
with torch.no_grad():
outputs = model(input_ids) outputs = model(input_ids)
prediction_logits = outputs.prediction_logits prediction_logits = outputs.prediction_logits
seq_relationship_logits = outputs.seq_relationship_logits seq_relationship_logits = outputs.seq_relationship_logits
@@ -920,6 +922,7 @@ class BigBirdModelIntegrationTest(unittest.TestCase):
model.eval() model.eval()
input_ids = torch.tensor([200 * [10] + 40 * [2] + [1]], device=torch_device, dtype=torch.long) input_ids = torch.tensor([200 * [10] + 40 * [2] + [1]], device=torch_device, dtype=torch.long)
with torch.no_grad():
output = model(input_ids).to_tuple()[0] output = model(input_ids).to_tuple()[0]
# fmt: off # fmt: off