Fix function signature

ext-torch/torch_binding.cpp

@@ -5,7 +5,7 @@
 
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
-  ops.def("mel_spectrogram(Tensor out, Tensor samples, Tensor filters, int64_t fft_size, int64_t fft_step) -> Tensor", &mel_spectrogram);
+  ops.def("mel_spectrogram(Tensor! out, Tensor samples, Tensor filters, int fft_size, int fft_step) -> ()", &mel_spectrogram);
   ops.impl("mel_spectrogram", torch::kCUDA, &mel_spectrogram);
 }
 

ext-torch/torch_binding.h

@@ -3,8 +3,8 @@
 #include <torch/torch.h>
 
 void mel_spectrogram(
-    torch::Tensor out,
-    torch::Tensor samples,
-    torch::Tensor filters,
+    torch::Tensor &out,
+    torch::Tensor const &samples,
+    torch::Tensor const &filters,
     int64_t fft_size,
-    int64_t fft_step);
+    int64_t fft_step);

mel_spectrogram/mel_spectrogram.cu

@@ -100,11 +100,11 @@ __global__ void normalize_mel_spec_kernel(
 }
 
 void mel_spectrogram(
-    torch::Tensor& out,        // [n_mel, n_frames]
-    torch::Tensor& samples,    // [n_samples]
-    torch::Tensor& filters,    // [n_mel, n_fft]
-    const int fft_size,
-    const int fft_step) {
+    torch::Tensor &out,        // [n_mel, n_frames]
+    torch::Tensor const &samples,    // [n_samples]
+    torch::Tensor const &filters,    // [n_mel, n_fft]
+    const int64_t fft_size,
+    const int64_t fft_step) {
     
     TORCH_CHECK(samples.is_cuda(), "Input samples must be a CUDA tensor");
     TORCH_CHECK(filters.is_cuda(), "Mel filters must be a CUDA tensor");
@@ -160,4 +160,4 @@ void mel_spectrogram(
         n_frames);
     
     cufftDestroy(plan);
-}
+}