Skip to content

Commit d982512

Browse files
remove .item from subtensor Op
1 parent d9cb188 commit d982512

File tree

1 file changed

+5
-12
lines changed

1 file changed

+5
-12
lines changed

pytensor/link/pytorch/dispatch/subtensor.py

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,12 @@ def subtensor(x, *ilists):
2020
new_indices = []
2121

2222
for i in indices:
23-
new_start, new_step, new_stop = None, None, None
2423
if isinstance(i, slice):
25-
if i.start:
26-
new_start = i.start.item()
27-
if i.step:
28-
new_step = i.step.item()
29-
if new_step < 0:
30-
raise NotImplementedError(
31-
"Negative step sizes are not supported in Pytorch"
32-
)
33-
if i.stop:
34-
new_stop = i.stop.item()
35-
new_indices.append(slice(new_start, new_stop, new_step))
24+
if i.step and i.step < 0:
25+
raise NotImplementedError(
26+
"Negative step sizes are not supported in Pytorch"
27+
)
28+
new_indices.append(i)
3629
else:
3730
new_indices.append(i.tolist())
3831

0 commit comments

Comments
 (0)