Doing the following is working for me but I am still a bit unsure if the autocast state if being sent to multiple devices. @ptrblck could you advise?
for (i, (x, y)) in enumerate(trainLoader):
	with torch.cuda.amp.autocast(enabled=True):
		# send the input to the device
		(x, y) = (x.to(DEVICE), y.to(DEVICE))

		# perform a forward pass and calculate the training loss
		pred = model(x)
		loss = lossFunc(pred, y)

	# calculate the gradients
	scaler.scale(loss).backward()
	scaler.step(opt)
	scaler.update()
	opt.zero_grad()