it probably should be just Data

pyg-team · Oct 14, 2022 · 63715bb · 63715bb
1 parent 345f66c
commit 63715bb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_geometric/transforms/random_shear.py b/torch_geometric/transforms/random_shear.py
@@ -2,7 +2,7 @@
 
 import torch
 
-from torch_geometric.data import Data, HeteroData
+from torch_geometric.data import Data
 from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform, LinearTransformation
 
@@ -29,7 +29,7 @@ class RandomShear(BaseTransform):
     def __init__(self, shear: Union[float, int]):
         self.shear = abs(shear)
 
-    def __call__(self, data: Union[Data, HeteroData]) -> Data:
+    def __call__(self, data: Data) -> Data:
         dim = data.pos.size(-1)
 
         matrix = data.pos.new_empty(dim, dim).uniform_(-self.shear, self.shear)