Hello,
I am fairly new to pytorch, so I’m not sure I’m doing things correctly…
I am using pytorch = 1.7.1
I have the following code for training/testing a Unet, I’ve made sure to set everything to the device that is set to CUDA, I even checked it by print(var.is_cuda). However, for some reason, during the testing part, my GPU usage drops to 0 and my CPU goes to 100%, though everything is set to CUDA.
I am looking for any explanation and help to fix it 
def train_model(model,
				device,
				crit="mse",
				opti="adam",
				dataset_file: str = "dict_wds-1201-0331-dec-march",
				num_epochs: int = 176,
				batch_size: int = 1,
				learning_rate: float = 1e-3,
				gen_load: bool = False,
				checkpoint_save: bool = True,
				):


	train_loader, test_loader, train_dataset, test_dataset = load_data(batch_size, dataset_file=dataset_file, discretize=discretize)
	n_train = len(train_dataset)
	n_test = len(test_dataset)

	writer = SummaryWriter(comment=f'LR_{learning_rate}_BS_{batch_size}_E_{num_epochs}_C_{crit}_O_{opti}_L_{n_train}_D_{discretize}')
	global_step = 0
	ets_benckmark = 0.0

	print("params #: ", sum(x.numel() for x in model.parameters()))

	criterion = nn.MSELoss()

	criterion = criterion.to(device)

	optimizer = torch.optim.Adam(
		model.parameters(),
		lr=learning_rate,
		weight_decay=1e-8
	)

	scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 2)

	for epoch in range(num_epochs):
		running_loss = 0.0
		loss_values = []
               
               # TRAINING PART

		with tqdm(train_loader, desc=f"Epoch [{epoch + 1}/{num_epochs}]", colour="#009B77") as t:
			for i, (ecmwf, truth) in enumerate(t):
				model.train()

				ecmwf = ecmwf.to(device, dtype=torch.float32)
				truth = truth.to(device, dtype=torch.float32)

				optimizer.zero_grad()
				outputs = model(ecmwf)

				if crit == "rmse":
					loss = torch.sqrt(criterion(outputs, truth))
				else:
					loss = criterion(outputs, truth)
				running_loss += loss.item()

				loss_values.append(loss.item())
				loss.backward()

				nn.utils.clip_grad_value_(model.parameters(), 0.1)
				optimizer.step()

				t.set_postfix(**{"iter": f"{i + 1}/{n_train // batch_size}", "batch loss": loss.item(),
								"avg loss": sum(loss_values)/n_train})
				writer.add_scalar("Loss/Train", loss.item(), global_step)

				global_step += 1


                                # TESTING PART

				if global_step % (n_train // (batch_size)) == 0:
					tot_l1_loss = 0
					tot_mse = 0
					tot_ets = 0
					with tqdm(test_loader, desc=f"Test", colour="#D2691E") as tl:
						for k, (ecmwf, truth) in enumerate(tl):
							for name, param in model.named_parameters():

								name = name.replace(".", "/")
								writer.add_histogram("Weights/" + name, param.data.cpu().numpy(), global_step)
								writer.add_histogram("Gradients/" + name, param.grad.cpu().numpy(), global_step)

							model.eval()
							with torch.no_grad():
								ecmwf = ecmwf.to(device)
								truth = truth.to(device)

								outputs = model(ecmwf)

								tot_mse += F.mse_loss(truth, outputs)

							# print(f"Validation mse: {tot_mse / len(test_loader)}")
							writer.add_scalar('mse/test', tot_mse / len(test_loader), global_step)

							tl.set_postfix(**{
								"Validation mse": f"{tot_mse / len(test_loader)}",
							})

							if k == len(test_loader)-1:
								writer.add_images('truth', truth, global_step)
								writer.add_images('pred', outputs, global_step)

					
					val_loss = tot_mse / len(test_loader)

		scheduler.step(val_loss)

		print(f"ETS benchmark avg on train_set (ecmwf/cmorph): {ets_benckmark / (n_train)}")

		if checkpoint_save:
			torch.save(model.state_dict(), f"./cp.pth")
			print(f"Checkpoint reached! :)")

	writer.close()

def main():
	args = arg_parse()

	device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
	torch.cuda.empty_cache()

	if device.type == "cuda":
		torch.backends.cudnn.benchmark = True

	model = my_UNet()

	model.to(device=device)

	train_model()