AnBucquet:

CUDA_LAUNCH_BLOCKING=1


Hi,
Thank you for your answer, and sorry for the delay. I am working from a jupiter notebook on floydhub so I set CUDA_LAUNCH_BLOCKING=1 into the terminal before running the notebook (I guess that is how to do it?). I have a different error:
RuntimeError                              Traceback (most recent call last)
 in 
11         l = torch.nn.functional.cross_entropy(prediction, y_train_batch)
12         my_optimizer.zero_grad()
—> 13         l.backward()
14         my_optimizer.step()
15         train_loss += l
/usr/local/lib/python3.6/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
91                 products. Defaults to False.
92         “”"
—> 93         torch.autograd.backward(self, gradient, retain_graph, create_graph)
94
95     def register_hook(self, hook):
/usr/local/lib/python3.6/site-packages/torch/autograd/init.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
88     Variable._execution_engine.run_backward(
89         tensors, grad_tensors, retain_graph, create_graph,
—> 90         allow_unreachable=True)  # allow_unreachable flag
91
92
RuntimeError: CuDNN error: CUDNN_STATUS_EXECUTION_FAILED
Also, to explain what I do: I try to implement in PyTorch a segmentation protocol with U-Net that was existing in Tensorflow. For the tensorflow code the segmentation results were provided as 4 binary mask images. Since in PyTorch the cross_entropy function does not allow multi-channel target, I multiply each image with a factor (1,100,200,255) and then summed them to obtain a single image in level of gray. Is it a correct way to proceed?

x_train_batch shape : torch.Size([10, 1, 256, 256])  dtype: torch.float32
y_train_batch shape: torch.Size([10, 256, 256]), dtype: torch.int64