diff --git a/spi4ba.c b/spi4ba.c
index a44e067..902f073 100644
--- a/spi4ba.c
+++ b/spi4ba.c
@@ -39,12 +39,17 @@
 /* Enter 4-bytes addressing mode (without sending WREN before) */
 int spi_enter_4ba_b7(struct flashctx *flash)
 {
+	int result;
 	const unsigned char cmd[JEDEC_ENTER_4_BYTE_ADDR_MODE_OUTSIZE] = { JEDEC_ENTER_4_BYTE_ADDR_MODE };
 
 	msg_trace("-> %s\n", __func__);
 
 	/* Switch to 4-bytes addressing mode */
-	return spi_send_command(flash, sizeof(cmd), 0, cmd, NULL);
+	result = spi_send_command(flash, sizeof(cmd), 0, cmd, NULL);
+	if (!result)
+		flash->in_4ba_mode = true;
+
+	return result;
 }
 
 /* Enter 4-bytes addressing mode with sending WREN before */
@@ -75,18 +80,25 @@
 	result = spi_send_multicommand(flash, cmds);
 	if (result)
 		msg_cerr("%s failed during command execution\n", __func__);
+	else
+		flash->in_4ba_mode = true;
 	return result;
 }
 
 /* Exit 4-bytes addressing mode (without sending WREN before) */
 int spi_exit_4ba_e9(struct flashctx *flash)
 {
+	int result;
 	const unsigned char cmd[JEDEC_EXIT_4_BYTE_ADDR_MODE_OUTSIZE] = { JEDEC_EXIT_4_BYTE_ADDR_MODE };
 
 	msg_trace("-> %s\n", __func__);
 
 	/* Switch to 3-bytes addressing mode  */
-	return spi_send_command(flash, sizeof(cmd), 0, cmd, NULL);
+	result = spi_send_command(flash, sizeof(cmd), 0, cmd, NULL);
+	if (!result)
+		flash->in_4ba_mode = false;
+
+	return result;
 }
 
 /* Exit 4-bytes addressing mode with sending WREN before */
@@ -115,9 +127,10 @@
 
 	/* Switch to 3-bytes addressing mode  */
 	result = spi_send_multicommand(flash, cmds);
-	if (result) {
+	if (result)
 		msg_cerr("%s failed during command execution\n", __func__);
-	}
+	else
+		flash->in_4ba_mode = false;
 	return result;
 }