diff --git a/short_bch_code_encoder.hh b/short_bch_code_encoder.hh
index 173f6e1..30bc7c7 100644
--- a/short_bch_code_encoder.hh
+++ b/short_bch_code_encoder.hh
@@ -34,11 +34,12 @@ public:
 	{
 		return (msg << P) | par[msg];
 	}
-	void operator()(int8_t *code, int msg)
+	int operator()(int8_t *code, int msg)
 	{
 		int cw = (*this)(msg);
 		for (int i = 0; i < N; ++i)
 			code[i] = 1 - 2 * ((cw >> i) & 1);
+		return cw;
 	}
 };
 
diff --git a/tests/soft_bch_regression_test.cc b/tests/soft_bch_regression_test.cc
index e6d2cdb..f47b8ee 100644
--- a/tests/soft_bch_regression_test.cc
+++ b/tests/soft_bch_regression_test.cc
@@ -64,7 +64,7 @@ int main()
 			float symb[CODE_LEN];
 
 			int dat = data();
-			encode(code, dat);
+			int enc = encode(code, dat);
 
 			for (int i = 0; i < CODE_LEN; ++i)
 				orig[i] = code[i];
@@ -85,17 +85,17 @@ int main()
 			for (int i = 0; i < CODE_LEN; ++i)
 				noisy[i] = code[i];
 
-			int dec = decode(code) >> (CODE_LEN - DATA_LEN);
+			int dec = decode(code);
 
 			for (int i = 0; i < CODE_LEN; ++i)
 				awgn_errors += noisy[i] * orig[i] < 0;
 			for (int i = 0; i < CODE_LEN; ++i)
 				quantization_erasures += !noisy[i];
-			uncorrected_errors += popcnt(dat^dec);
-			for (int i = 0; i < DATA_LEN; ++i)
-				decoder_errors += ((dec^dat)&(1<<i)) && orig[i] * noisy[i] > 0;
+			uncorrected_errors += popcnt(enc^dec);
+			for (int i = 0; i < CODE_LEN; ++i)
+				decoder_errors += ((enc^dec)&(1<<i)) && orig[i] * noisy[i] > 0;
 		}
-		float bit_error_rate = (float)uncorrected_errors / (float)(DATA_LEN * LOOPS);
+		float bit_error_rate = (float)uncorrected_errors / (float)(CODE_LEN * LOOPS);
 		if (bit_error_rate < 0.0001)
 			min_SNR = std::min(min_SNR, SNR);