@@ -773,7 +773,14 @@ class LlamaIntegrationTest(unittest.TestCase):
|
||||
8: torch.tensor([[-6.5208, -4.1218, -4.9377, -3.2536, 0.8127, -2.9811, 1.2918, -3.3848]])
|
||||
}
|
||||
|
||||
self.assertTrue(torch.allclose(EXPECTED_MEAN[self.cuda_compute_capability_major_version].to(torch_device), out.logits.mean(-1), atol=1e-2, rtol=1e-2))
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
EXPECTED_MEAN[self.cuda_compute_capability_major_version].to(torch_device),
|
||||
out.logits.float().mean(-1),
|
||||
atol=1e-2,
|
||||
rtol=1e-2
|
||||
)
|
||||
)
|
||||
|
||||
# slicing logits[0, 0, 0:15]
|
||||
EXPECTED_SLICE = {
|
||||
@@ -785,7 +792,7 @@ class LlamaIntegrationTest(unittest.TestCase):
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
EXPECTED_SLICE[self.cuda_compute_capability_major_version].to(torch_device),
|
||||
out.logits[0, 0, :15],
|
||||
out.logits[0, 0, :15].float(),
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
)
|
||||
@@ -810,7 +817,14 @@ class LlamaIntegrationTest(unittest.TestCase):
|
||||
8: torch.tensor([[-6.6544, -4.1259, -4.9840, -3.2456, 0.8261, -3.0124, 1.2971, -3.3641]])
|
||||
}
|
||||
|
||||
self.assertTrue(torch.allclose(EXPECTED_MEAN[self.cuda_compute_capability_major_version].to(torch_device), out.logits.mean(-1), atol=1e-2, rtol=1e-2))
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
EXPECTED_MEAN[self.cuda_compute_capability_major_version].to(torch_device),
|
||||
out.logits.float().mean(-1),
|
||||
atol=1e-2,
|
||||
rtol=1e-2
|
||||
)
|
||||
)
|
||||
|
||||
# slicing logits[0, 0, 0:15]
|
||||
EXPECTED_SLICE = {
|
||||
@@ -822,7 +836,7 @@ class LlamaIntegrationTest(unittest.TestCase):
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
EXPECTED_SLICE[self.cuda_compute_capability_major_version].to(torch_device),
|
||||
out.logits[0, 0, :15],
|
||||
out.logits[0, 0, :15].float(),
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user