add logic to allow negative length slicing

finegrain-ai · Dec 12, 2023 · 06822af · 06822af
1 parent 7992258
commit 06822af
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/refiners/fluxion/layers/basics.py b/src/refiners/fluxion/layers/basics.py
@@ -91,6 +91,8 @@ def __init__(self, dim: int, start: int, length: int) -> None:
         self.length = length
 
     def forward(self, x: Tensor) -> Tensor:
+        if self.length < 0:
+            return x.narrow(self.dim, self.start, x.shape[self.dim] - self.start + self.length + 1)
         return x.narrow(self.dim, self.start, self.length)