diff --git a/comfy_latent_upscaler.py b/comfy_latent_upscaler.py
index 605d812..081ad1c 100644
--- a/comfy_latent_upscaler.py
+++ b/comfy_latent_upscaler.py
@@ -9,7 +9,7 @@ class Upscaler(nn.Module):
 		Basic NN layout, ported from:
 		https://github.com/city96/SD-Latent-Upscaler/blob/main/upscaler.py
 	"""
-	version = 2.0 # network revision
+	version = 2.1 # network revision
 	def head(self):
 		return [
 			nn.Conv2d(self.chan, self.size, kernel_size=self.krn, padding=self.pad),
diff --git a/train.py b/train.py
index 8381f36..3d6b681 100644
--- a/train.py
+++ b/train.py
@@ -112,8 +112,8 @@ if __name__ == "__main__":
 		eval_src = torch.from_numpy(np.load(f"test_{args.ver}_{args.res}px.npy")).to(target_dev)
 		eval_dst = torch.from_numpy(np.load(f"test_{args.ver}_{dst_res}px.npy")).to(target_dev)
 	else:
-		eval_src = dataset[0][0]
-		eval_dst = dataset[0][1]
+		eval_src = torch.unsqueeze(dataset[0][0],0)
+		eval_dst = torch.unsqueeze(dataset[0][1],0)
 
 	model = Upscaler(args.fac)
 	if args.resume:
@@ -131,6 +131,7 @@ if __name__ == "__main__":
 		total_steps=int(args.steps/args.bs),
 		max_lr=float(args.lr)/args.bs,
 		pct_start=0.015,
+		final_div_factor=2500,
 	)
 	# scaler = torch.cuda.amp.GradScaler()
 	progress = tqdm(total=args.steps)