fix rope unit tests

ebsmothers · ebsmothers · commit 78e254c646e6 · 2024-08-10T11:13:51.000-07:00
diff --git a/tests/torchtune/models/llama3_1/test_position_embeddings.py b/tests/torchtune/models/llama3_1/test_position_embeddings.py
@@ -132,7 +132,7 @@ def test_rope_init_meta_device(self, input_params):
         with torch.device("meta"):
             meta_rope = Llama3ScaledRoPE(dim=head_dim, max_seq_len=max_seq_len)
 
-        meta_rope._rope_init()
+        meta_rope.rope_init()
         for p1, p2 in zip(rope_on_device.buffers(), meta_rope.buffers()):
             torch.testing.assert_close(p1, p2)
 
diff --git a/tests/torchtune/modules/test_position_embeddings.py b/tests/torchtune/modules/test_position_embeddings.py
@@ -128,7 +128,7 @@ def test_rope_init_meta_device(self, input_params):
                 dim=head_dim, max_seq_len=max_seq_len
             )
 
-        meta_rope._rope_init()
+        meta_rope.rope_init()
         for p1, p2 in zip(rope_on_device.buffers(), meta_rope.buffers()):
             torch.testing.assert_close(p1, p2)
 

Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,7 @@ def test_rope_init_meta_device(self, input_params):`
`128`	`128`	`dim=head_dim, max_seq_len=max_seq_len`
`129`	`129`	`)`
`130`	`130`
`131`		`- meta_rope._rope_init()`
	`131`	`+ meta_rope.rope_init()`
`132`	`132`	`for p1, p2 in zip(rope_on_device.buffers(), meta_rope.buffers()):`
`133`	`133`	`torch.testing.assert_close(p1, p2)`
`134`	`134`