slice_by_index operation not supported on GPU, defaulting to CPU
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## 🐞Describing the bug
Running into an issue where my entire model is able to compile via ct.convert and run on GPU except for the slice_by_index operations:

Interestingly, this only occurs when I run on my iPhone on device and does not when I run on my M4 Macbook Pro, which assigns whole model onto GPU successfully.
The reason I'm asking about this here is because I wonder if there's a way to do the same operation without necessarily slicing or if there's some core code I can edit in the package to fix this. Tried things like calling z.contiguous() prior to slicing, adding static slicing where indices are int constants determined by if statements. Also tried F.interpolate but doesn't result in the same output, and torch.narrow calls mb.slice_by_index() so results in no changes.
## To Reproduce
The issue arises from `z = z[..., self.pad:-self.pad, :]` and `z = z[..., self.pad:self.pad + length]`, which are just [2:-2], [2:-3], [2:-4] dependent on the layer number.
```
class HDecLayerNew(nn.Module):
def __init__(self, chin, chout, last=False, kernel_size=8, stride=4, norm_groups=1, empty=False,
freq=True, dconv=True, norm=True, context=1, dconv_kw={}, pad=True,
context_freq=True, rewrite=True):
super().__init__()
self.pad = kernel_size // 4 if pad else 0
self.last = last
self.freq = freq
self.chin = chin
self.empty = empty
self.stride = stride
self.kernel_size = kernel_size
self.norm = norm
self.context_freq = context_freq
norm_fn = lambda d: nn.Identity() # noqa
if norm:
norm_fn = lambda d: nn.GroupNorm(norm_groups, d) # noqa
self.conv_tr = nn.ConvTranspose1d(chin, chout, kernel_size, stride)
self.norm2 = norm_fn(chout)
if not self.empty:
klass = nn.Conv2d if freq else nn.Conv1d
if rewrite:
self.rewrite = klass(chin, 2 * chin, 1 + 2 * context, 1, context) if context_freq else klass(chin, 2 * chin, [1, 1 + 2 * context], 1, [0, context])
self.norm1 = norm_fn(2 * chin)
else:
self.rewrite = None
self.dconv = DConv(chin, **dconv_kw) if dconv else None
def forward(self, x, skip, length):
if self.freq and x.dim() == 3:
B, C, T = x.shape
x = x.view(B, self.chin, -1, T)
if not self.empty:
x = x + skip
y = F.glu(self.norm1(self.rewrite(x)), dim=1) if self.rewrite else x
if self.dconv:
if self.freq:
B, C, Fr, T = y.shape
y = y.permute(0, 2, 1, 3).reshape(-1, C, T)
y = self.dconv(y)
if self.freq:
y = y.view(B, Fr, C, T).permute(0, 2, 1, 3)
else:
y = x
assert skip is None
if self.freq:
B, C, Fr, T = y.shape
y_reshaped = y.permute(0, 3, 1, 2).reshape(B * T, C, Fr)
z_reshaped = self.conv_tr(y_reshaped)
C_out, Fr_out = z_reshaped.shape[1], z_reshaped.shape[2]
z = z_reshaped.view(B, T, C_out, Fr_out).permute(0, 2, 3, 1)
z = self.norm2(z)
if self.pad:
z = z[..., self.pad:-self.pad, :]
if not self.last:
z = F.gelu(z)
else:
z = self.norm2(self.conv_tr(y))
if self.pad:
z = z[..., self.pad:self.pad + length]
assert z.shape[-1] == length, (z.shape[-1], length)
if not self.last:
z = F.gelu(z)
return z, y
```
## System environment (please complete the following information):
- coremltools version: 8.3.0
- Source dialect: torchscript (torch.jit.trace)
- OS: MacOS 15.3.2, iOS 18.5
- PyTorch version: 2.4.1
Contributor guide
Research direction
Start by reproducing the torch.jit.trace and ct.convert flow with the reported slice expressions, then compare execution on the iPhone GPU and M4 Mac GPU. Trace how slice_by_index is assigned during conversion and verify that the model preserves its output without falling back to CPU on iOS.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- ios, python, pytorch
- Domain
- machine-learning, mobile-dev
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100