/opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THC/THCTensorScatterGather.cu:100: void THCudaTensor_gatherKernel(TensorInfo<Real, IndexType>, TensorInfo<Real, IndexType>, TensorInfo<long, IndexType>, int, IndexType) [with IndexType = unsigned int, Real = float, Dims = 3]: block: [40,0,0], thread: [31,0,0] Assertion `indexValue >= 0 && indexValue < src.sizes[dim]` failed.
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THC/generic/THCTensorScatterGather.cu line=71 error=59 : device-side assert triggered
Traceback (most recent call last):
  File "cudafail.py", line 47, in <module>
    grad_sum.backward(retain_graph=False)
  File "/home/michael/miniconda2/envs/pt/lib/python3.6/site-packages/torch/tensor.py", line 107, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/michael/miniconda2/envs/pt/lib/python3.6/site-packages/torch/autograd/__init__.py", line 93, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: cuda runtime error (59) : device-side assert triggered at /opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THC/generic/THCTensorScatterGather.cu:71

import torch
from torch import nn
from torchvision import datasets, transforms

class Net(nn.Module):
	def __init__(self):
		super(Net, self).__init__()
		self.conv1 = nn.Conv2d(3, 20, kernel_size=5, bias=False)
		self.conv2 = nn.Conv2d(20, 40, kernel_size=5, bias=False)
		self.linear1 = nn.Linear(40 * 5 * 5, 300, bias=False)
		self.linear2 = nn.Linear(300, 10, bias=False)
		self.pool = nn.MaxPool2d(2, 2)
		self.relu = nn.ReLU()

	def forward(self, input):
		x = self.relu(self.pool(self.conv1(input)))
		x = self.relu(self.pool(self.conv2(x)))
		x = x.view(x.size(0), -1)
		x = self.relu(self.linear1(x))
		return self.linear2(x)

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=8)

model = Net().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1, momentum=0, nesterov=False)

for epoch in range(100):
        print(epoch)
	model.train()
	for i, (images, labels) in enumerate(trainloader, 0):
		images = images.cuda()
		labels = labels.cuda()
		outputs = model(images)
		loss = nn.CrossEntropyLoss()(outputs, labels)
		optimizer.zero_grad()
		loss.backward(retain_graph=True)

		grads = torch.autograd.grad(loss, model.parameters(), create_graph=True)
		grad_sum = 0
		for grad in grads:
			grad_sum += 0.1 * grad.pow(2).sum()
		grad_sum.backward(retain_graph=False)

		optimizer.step()