input must be a CUDA tensor error when using AdaPool3d

[Tloops]

Great work! However, when i'm using AdaPool3d, and i encoutered the error below:

Traceback (most recent call last):
  File "<masked>/main.py", line 339, in <module>
    main()
  File "<masked>/main.py", line 329, in main
    train_loss, train_acc = train(model, train_loader, epoch, criterion, optimizer)
  File "<masked>/main.py", line 113, in train
    outputs = model(inputs)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py", line 183, in forward
    return self.module(*inputs[0], **module_kwargs[0])
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "<masked>/models/cnn3d_adapool.py", line 52, in forward
    x = self.pool(x)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/adaPool-0.1-py3.9-linux-x86_64.egg/adaPool/idea.py", line 1495, in forward
    return adapool3d(x, beta=self.beta, kernel_size=self.kernel_size, stride=self.stride, return_mask=self.return_mask, native=self.native)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/adaPool-0.1-py3.9-linux-x86_64.egg/adaPool/idea.py", line 771, in adapool3d
    x = beta*CUDA_ADAPOOL3d_EDSCW.apply(x, kernel_size, stride, return_mask) + (1. - beta)*CUDA_ADAPOOL3d_EM.apply(x, kernel_size, stride, return_mask)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/autograd/function.py", line 553, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/torch/cuda/amp/autocast_mode.py", line 123, in decorate_fwd
    return fwd(*args, **kwargs)
  File "<masked>/miniconda3/envs/python39/lib/python3.9/site-packages/adaPool-0.1-py3.9-linux-x86_64.egg/adaPool/idea.py", line 505, in forward
    adapool_cuda.forward_3d_edscw(input.contiguous(), kernel, stride, output, return_mask, mask)
RuntimeError: input.is_cuda() INTERNAL ASSERT FAILED at "CUDA/adapool_cuda.cpp":616, please report a bug to PyTorch. input must be a CUDA tensor

I've checked the if the input is on the cuda by printing input.is_cuda just before the adapool_cuda.forward_3d_edscw function call, and it displays True. But, when it comes to the cpp file, the input became NOT a CUDA tensor. I'm really comfused about that. Hope to receive your reply soon. Thanks!