Repository: pytorch
Issue: RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1607370116979/work/aten/src/ATen/native/cuda/Loops.cuh":94, please report a bug to PyTorch. 
Labels: needs reproduction,module: cuda,triaged,module: TensorIterator

## 🐛 Bug

Hi, when I am running the code,  I met this bug: RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1607370116979/work/aten/src/ATen/native/cuda/Loops.cuh":94, please report a bug to PyTorch. 

## To Reproduce

Steps to reproduce the behavior:

1. my code is available [here](https://drive.google.com/file/d/19_Jeqa889ue3JRNYMIVyfrE3eHVWgwJu/view?usp=sharing). Just simply run `main.ipynb`.
2. I'm trying to reproduce a specific triplet learning algorithm from this paper: https://arxiv.org/abs/2007.11186. For training, the code runs well on the CPU, but when switched to GPU (both single/data-parallel don't work), this error occurs.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-9-cf6df666728d> in <module>()
      6         output1, output2, output3 = model(anchor, positive, negative)
      7         optimizer.zero_grad()
----> 8         loss = loss_fn(output1, output2, output3)
      9         loss.backward()
     10         optimizer.step()

/opt/anaconda3/envs/cryo-torch/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/root/cryo/code/self-supervised/triplet/losses.py in forward(self, anchor, positive, negative, size_average)
     18         Dap = F.pairwise_distance(anchor, positive, p=2).pow(2)
     19         Dan = F.pairwise_distance(anchor, negative, p=2).pow(2)
---> 20         Lst = torch.max(torch.tensor(0), Dap - Dan + self.margin1).squeeze()
     21 
     22         new_neg = negative.reshape([1, 128, 1, 1])

RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1607370116979/work/aten/src/ATen/native/cuda/Loops.cuh":94, please report a bug to PyTorch. 


## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
I wish to train the network on GPU as it works on CPU. It seems something wrong with the Loss function part I wrote? I am new to it, any help will be greatly appreciated.


## Environment

 PyTorch version: 1.7.1
Is debug build: False
CUDA used to build PyTorch: 10.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.3 LTS (x86_64)
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.6 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti
GPU 4: GeForce RTX 2080 Ti
GPU 5: GeForce RTX 2080 Ti
GPU 6: GeForce RTX 2080 Ti
GPU 7: GeForce RTX 2080 Ti

Nvidia driver version: 440.33.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.7.1
[pip3] torchvision==0.8.2
[conda] blas                      1.0                         mkl    defaults
[conda] cudatoolkit               10.1.243             h036e899_8    conda-forge
[conda] mkl                       2020.2                      256    defaults
[conda] mkl-service               2.3.0            py36h8c4c3a4_2    conda-forge
[conda] mkl_fft                   1.2.1            py36h54f3939_0    defaults
[conda] mkl_random                1.2.0            py36h7c3b610_1    conda-forge
[conda] numpy                     1.19.2           py36h54aff64_0    defaults
[conda] numpy-base                1.19.2           py36hfa32c7d_0    defaults
[conda] pytorch                   1.7.1           py3.6_cuda10.1.243_cudnn7.6.3_0    pytorch
[conda] torchvision               0.8.2                py36_cu101    pytorch

Thanks in advance!

cc @ngimel