Skip to content

Commit d34e87c

Browse files
bottlermeta-codesync[bot]
authored andcommitted
validate input and edges dtype in GatherScatter python wrapper
Summary: Add explicit dtype checks for input (torch.float32) and edges (torch.int64) in GatherScatter.forward and gather_scatter_python to match C++ TensorAccessor<float,2> and TensorAccessor<int64_t,2> expectations. Python previously validated ndim, shape, and input dtype in forward but not edges dtype, and gather_scatter_python lacked dtype checks entirely, relying on ATen error from accessor. This makes errors python-friendly and guards C++ accessor before TensorAccessor construction. ___ Differential Revision: D108140422 fbshipit-source-id: ba54e857279a480a02e2c8f27e316f2e23cc6092
1 parent 1f7f85c commit d34e87c

1 file changed

Lines changed: 6 additions & 0 deletions

File tree

pytorch3d/ops/graph_conv.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,10 @@ def gather_scatter_python(input, edges, directed: bool = False):
114114
raise ValueError("edges can only have 2 dimensions.")
115115
if not (edges.shape[1] == 2):
116116
raise ValueError("edges must be of shape (num_edges, 2).")
117+
if not (input.dtype == torch.float32):
118+
raise ValueError("input has to be of type torch.float32.")
119+
if not (edges.dtype == torch.int64):
120+
raise ValueError("edges has to be of type torch.int64.")
117121

118122
num_vertices, input_feature_dim = input.shape
119123
num_edges = edges.shape[0]
@@ -152,6 +156,8 @@ def forward(ctx, input, edges, directed=False):
152156
raise ValueError("edges must be of shape (num_edges, 2).")
153157
if not (input.dtype == torch.float32):
154158
raise ValueError("input has to be of type torch.float32.")
159+
if not (edges.dtype == torch.int64):
160+
raise ValueError("edges has to be of type torch.int64.")
155161

156162
ctx.directed = directed
157163
input, edges = input.contiguous(), edges.contiguous()

0 commit comments

Comments
 (0)