Nx can’t do dynamic shape modifications. You have a few options here:
- Crop with Evision
- Figure out the slice sizes outside of defn (by using Nx.argmin to figure out the starts and Nx.sum to figure out the lengths) and then use Nx.slice with those now-fixed values.