just use tuple

geoopt · Jun 19, 2019 · 7b09e05 · 7b09e05
1 parent 289c705
commit 7b09e05
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/geoopt/manifolds/poincare/__init__.py b/geoopt/manifolds/poincare/__init__.py
@@ -97,11 +97,10 @@ def transp(self, x, y, v, *more, dim=-1):
         if not more:
             return math.parallel_transport(x, y, v, c=self.c, dim=dim)
         else:
-            vecs = torch.stack((v,) + more, dim=0)
-            transp = math.parallel_transport(
-                x, y, vecs, c=self.c, dim=idx2sign(dim, x.dim())
+            return tuple(
+                math.parallel_transport(x, y, vec, c=self.c, dim=dim)
+                for vec in (v, *more)
             )
-            return transp.unbind(0)
 
     def transp_follow_retr(self, x, u, v, *more, dim=-1):
         y = self.retr(x, u, dim=dim)