Fixing conversation test for torch 1.8 (#10545)

54e55b52 · Nicolas Patry · GitHub · dc9aaa38 · 54e55b52
Unverified Commit 54e55b52 authored Mar 05, 2021 by Nicolas Patry Committed by GitHub Mar 05, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

tests/test_pipelines_conversational.py tests/test_pipelines_conversational.py +2 -2

No files found.
--- a/tests/test_pipelines_conversational.py
+++ b/tests/test_pipelines_conversational.py
@@ -53,9 +53,9 @@ class SimpleConversationPipelineTests(unittest.TestCase):
        model = GPT2LMHeadModel(config)
        # Force model output to be L
        V, D = model.lm_head.weight.shape
-        bias = torch.zeros(V, requires_grad=True)
-        weight = torch.zeros((V, D), requires_grad=True)
+        bias = torch.zeros(V)
        bias[76] = 1
+        weight = torch.zeros((V, D), requires_grad=True)

        model.lm_head.bias = torch.nn.Parameter(bias)
        model.lm_head.weight = torch.nn.Parameter(weight)