diff --git a/aimodel/src/lib/ai/RainfallWaterContraster.py b/aimodel/src/lib/ai/RainfallWaterContraster.py
index a6d7890..2ab91fe 100644
--- a/aimodel/src/lib/ai/RainfallWaterContraster.py
+++ b/aimodel/src/lib/ai/RainfallWaterContraster.py
@@ -88,18 +88,15 @@ class RainfallWaterContraster(object):
 		i_batch = -1
 		for batch in dataset:
 			i_batch += 1
-			result_batch = self.model_predict(batch[0]) # ((rainfall, water), dummy_label)
-			rainfall, water = tf.unstack(result_batch, axis=-2)
+			rainfall = self.model_predict(batch) # ((rainfall, water), dummy_label)
 			
-			rainfall = tf.unstack(rainfall, axis=0)
-			water = tf.unstack(water, axis=0)
-			for step in zip(rainfall, water):
+			for step in tf.unstack(rainfall, axis=0):
 				yield step
 		
 	
-	def embed_rainfall(self, dataset):
-		result = []
-		for batch in dataset:
-			result_batch = self.model_predict(batch)
-			result.extend(tf.unstack(result_batch, axis=0))
-		return result
\ No newline at end of file
+	# def embed_rainfall(self, dataset):
+	# 	result = []
+	# 	for batch in dataset:
+	# 		result_batch = self.model_predict(batch)
+	# 		result.extend(tf.unstack(result_batch, axis=0))
+	# 	return result
\ No newline at end of file
diff --git a/aimodel/src/subcommands/pretrain_predict.py b/aimodel/src/subcommands/pretrain_predict.py
index 29cfecf..3298728 100644
--- a/aimodel/src/subcommands/pretrain_predict.py
+++ b/aimodel/src/subcommands/pretrain_predict.py
@@ -68,11 +68,8 @@ def run(args):
 		handle = handle_open(filepath_output, "w")
 	
 	i = 0
-	for rainfall, water in ai.embed(dataset):
-		handle.write(json.dumps({
-			"rainfall": rainfall.numpy().tolist(),
-			"water": water.numpy().tolist()
-		}, separators=(',', ':'))+"\n") # Ref https://stackoverflow.com/a/64710892/1460422
+	for rainfall in ai.embed(dataset):
+		handle.write(json.dumps(rainfall.numpy().tolist(), separators=(',', ':'))+"\n") # Ref https://stackoverflow.com/a/64710892/1460422
 		
 		if i == 0 or i % 1000 == 0:
 			sys.stderr.write(f"[pretrain:predict] STEP {i}\r")