Repository: pytorch
Issue: Numpy/scipy module works fine with Torch modules, but not TorchScript. How to torchscript a numpy/scipy module?
Labels: oncall: jit

### 🐛 Numpy module works fine with Torch modules, but not TorchScript.

```python
from scipy.signal import find_peaks

batch_size = 1
input_data_shape = 1000
input_shape = (batch_size, input_data_shape)

reference_inputs = numpy.random.random(input_shape)
reference_outputs, _ = find_peaks(reference_inputs[0, :])

class FindPeaks(torch.nn.Module):
    def __init__(self):
        super(FindPeaks, self).__init__()

    def forward(self, xs):
        xs_numpy = xs.numpy()[0, :]
        peaks, _ = find_peaks(xs_numpy)
        return torch.tensor(peaks, dtype=int)

inputs = torch.tensor(reference_inputs, dtype=float)
torch_model = FindPeaks()
torch_outputs = torch_model(inputs)

torchscript_model = torch.jit.trace(torch_model, example_inputs=[inputs])
torchscript_model.save(f"./artifacts/{torch_model.__class__.__name__}.pt")

torchscript_outputs = torchscript_model(inputs).detach()
assert isinstance(torchscript_outputs, torch.Tensor)
assert torchscript_outputs.shape == reference_outputs.shape
assert numpy.allclose(
    reference_outputs, torchscript_outputs.numpy(), rtol=1.0e-3, atol=1.0e-5
)

for i in range(5):
    reference_inputs = numpy.random.random(input_shape)
    reference_outputs, _ = find_peaks(reference_inputs[0, :])

    inputs = torch.tensor(reference_inputs, dtype=float)

    torch_outputs = torch_model(inputs).detach()
    assert isinstance(torch_outputs, torch.Tensor)
    assert torch_outputs.shape == reference_outputs.shape # works fine
    assert numpy.allclose(
        reference_outputs, torch_outputs.numpy(), rtol=1.0e-3, atol=1.0e-5
    ) # works fine

    torchscript_outputs = torchscript_model(inputs).detach()
    assert isinstance(torchscript_outputs, torch.Tensor)
    assert torchscript_outputs.shape == reference_outputs.shape, \
        (torchscript_outputs, reference_outputs) # not working, seems memorizing the input/output when compiling the model.
    assert numpy.allclose(
        reference_outputs, torchscript_outputs.numpy(), rtol=1.0e-3, atol=1.0e-5
    )
```

### Versions

```
Collecting environment information...
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.3 (x86_64)
GCC version: Could not collect
Clang version: 16.0.3
CMake version: version 3.21.1
Libc version: N/A

Python version: 3.8.16 (default, Dec  7 2022, 01:39:17)  [Clang 14.0.0 (clang-1400.0.29.202)] (64-bit runtime)
Python platform: macOS-13.3-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Intel(R) Core(TM) i9-9980HK CPU @ 2.40GHz

Versions of relevant libraries:
[pip3] flake8==4.0.1
[pip3] mypy==0.910
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.0
[pip3] torch==1.12.1
[pip3] torchaudio==0.12.1
[pip3] torchvision==0.13.1
[conda] Could not collect
```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel