Re-order reshape and device upload

tenstorrent · Dec 11, 2024 · 4e78614 · 4e78614
1 parent 8e49222
commit 4e78614
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/tt_eager/ops/test_tilize_zero_padding_channels_last.cpp b/tests/tt_eager/ops/test_tilize_zero_padding_channels_last.cpp
@@ -37,8 +37,9 @@ int main(int argc, char** argv) {
         ////////////////////////////////////////////////////////////////////////////
         ttnn::SimpleShape shape{1, 32, 61, 32};
         // Allocates a DRAM buffer on device populated with values specified by initialize
-        Tensor a = ttnn::arange(/*start=*/0, /*stop=*/shape.volume(), /*step=*/1, DataType::BFLOAT16, std::ref(*device))
-                       .reshape(shape);
+        Tensor a = ttnn::arange(/*start=*/0, /*stop=*/shape.volume(), /*step=*/1, DataType::BFLOAT16)
+                       .reshape(shape)
+                       .to(device);
         Tensor b = ttnn::tilize_with_zero_padding(a);
         Tensor c = b.cpu();
         ////////////////////////////////////////////////////////////////////////////