Create PyTorch script for gather

crates/burn-import/onnx-tests/tests/gather/gather.onnx

@@ -1,16 +1,16 @@
-pytorch2.2.2:�

-a
-onnx::GatherElements_0
-onnx::GatherElements_1
2/GatherElements"GatherElements*
+pytorch2.1.1:�

+A
+onnx::Gather_0
+onnx::Gather_1
2/Gather"Gather*
 axis
�

-main_graphZ(
-onnx::GatherElements_0
+main_graphZ 
+onnx::Gather_0
 

 
-Z(
-onnx::GatherElements_1
-
-
+Z
+onnx::Gather_1
+
+
 b
 
2
 


crates/burn-import/onnx-tests/tests/gather/gather.py

@@ -11,8 +11,8 @@ def __init__(self):
         super(Model, self).__init__()
 
     def forward(self, x, index):
-        x = torch.gather(x, 1, index)
-        return x
+        gathered = torch.index_select(x, 1, index)
+        return gathered
 
 
 def main():
@@ -24,19 +24,19 @@ def main():
     model.eval()
     device = torch.device("cpu")
     onnx_name = "gather.onnx"
-    dummy_input = torch.randn(2, 2, device=device)
-    dummy_index = torch.randint(high=2, size=(2, 2), device=device, dtype=torch.int64)
+
+    dummy_input = torch.randn(2, 3, device=device)
+    dummy_index = torch.tensor([0, 2], device=device, dtype=torch.int64)
 
     torch.onnx.export(model, (dummy_input, dummy_index), onnx_name,
                       verbose=False, opset_version=16)
 
     print("Finished exporting model to {}".format(onnx_name))
 
     # Output some test data for use in the test
-    test_input = torch.tensor([[1.0, 2.0],
-                               [3.0, 4.0]])
-    test_index = torch.tensor([[0, 0],
-                               [1, 0]])
+    test_input = torch.tensor([[1.0, 2.0, 3.0],
+                               [4.0, 5.0, 6.0]])
+    test_index = torch.tensor([0, 2], dtype=torch.int64)
 
     print("Test input data: {}, {}".format(test_input, test_index))
     output = model.forward(test_input, test_index)