Add no_grad context during prediction. (21d7b7bf) · Commits · Agrawal, Devanshu / bottleneck_nngp

torch_experiments/main.py

+6 −4

Original line number	Diff line number	Diff line
		@@ -69,6 +69,7 @@ for i in range(args.iters):
		print("iter {} loss: {:.3f}".format(i, loss.item()))

		# compute loss of final iteration
		with torch.no_grad():
		Ks = model(X, num_samples=args.train_samples)
		loss = -model.log_likelihood(Ks, Y)
		results["train"]["loss"].append(loss.item())
		@@ -79,6 +80,7 @@ results["train"]["time"] = time.time()-time_0

		# get final values with higher number of samples
		time_0 = time.time()
		with torch.no_grad():
		Ks = model(X, num_samples=args.test_samples)
		loss = -model.log_likelihood(Ks, Y)
		results["test"] = {}