diff --git a/README.md b/README.md
index 8f9bfa8..221773c 100644
--- a/README.md
+++ b/README.md
@@ -270,6 +270,7 @@ A third method provides flexibility in case the user wants to perform read or wr
         size: Optional[Union[int, Sequence[int]]] = None,
         pip: Optional[bool],
         verbose: Optional[bool] = False,
+        format_amba: Optional[bool] = False,
     ) -> Sequence[dict]:
 ```
 
@@ -281,6 +282,7 @@ A third method provides flexibility in case the user wants to perform read or wr
 * Optional[pip] - Define if the address/data phase will overlap in a pipeline manner or not, default non-pipelined 
 * Optional[sync] - Drive signals on next clk edge
 * Optional[verbose] - Print a msg on every txn
+* Optional[format_amba] - Enforce AMBA data masking/shifting according to spec (Table 6-1 Active byte lanes for a 32-bit little-endian data bus - ARM IHI 0033B.b)
 
 **Return**
 * Sequence[dict] - Once all transactions are dispatched and their responses are received, the function returns a list of dict with [AHBResp](cocotbext/ahb/ahb_types.py) responses along with the data.
diff --git a/cocotbext/ahb/ahb_master.py b/cocotbext/ahb/ahb_master.py
index 2a79568..d1d18ae 100644
--- a/cocotbext/ahb/ahb_master.py
+++ b/cocotbext/ahb/ahb_master.py
@@ -4,7 +4,7 @@
 # License           : MIT license <Check LICENSE>
 # Author            : Anderson I. da Silva (aignacio) <anderson@aignacio.com>
 # Date              : 08.10.2023
-# Last Modified Date: 24.10.2024
+# Last Modified Date: 27.12.2024
 import logging
 import cocotb
 import copy
@@ -93,7 +93,9 @@ def _check_size(size: int, data_bus_width: int) -> None:
         elif size <= 0 or (size & (size - 1)) != 0:
             raise ValueError(f"Error -> {size} - Size must" f"be a positive power of 2")
 
-    def _fmt_amba(self, address: Sequence[int], size: Sequence[int], value: Sequence[int]) -> Sequence[int]:
+    def _fmt_amba(
+        self, address: Sequence[int], size: Sequence[int], value: Sequence[int]
+    ) -> Sequence[int]:
         """Format the write data to follow AMBA by shifting / masking data."""
         new_val = []
         offset = (self.bus.data_width // 8) - 1
@@ -106,7 +108,7 @@ def _fmt_amba(self, address: Sequence[int], size: Sequence[int], value: Sequence
                     data = (val & 0xFFFF) << ((addr & offset) * 8)
                 elif sz == 1:
                     data = (val & 0xFF) << ((addr & offset) * 8)
-                new_val.append(data)
+                new_val.append(data & (2**self.bus.data_width - 1))
             else:
                 new_val.append(val)
         return new_val
@@ -313,9 +315,6 @@ async def write(
         if not isinstance(value, list):
             value = [value]
 
-        if format_amba is True:
-            value = self._fmt_amba(address, size, value)
-
         # if not isinstance(size, list):
         # size = [size]
 
@@ -332,6 +331,9 @@ async def write(
                 f"different from size length ({len(size)})"
             )
 
+        if format_amba is True:
+            value = self._fmt_amba(address, size, value)
+
         # Need to copy data as we'll have to shift address/value
         t_address = copy.deepcopy(address)
         t_value = copy.deepcopy(value)
@@ -417,6 +419,7 @@ async def custom(
         pip: Optional[bool] = True,
         verbose: Optional[bool] = False,
         sync: Optional[bool] = False,
+        format_amba: Optional[bool] = False,
     ) -> Sequence[dict]:
         """Back-to-Back operation"""
 
@@ -447,6 +450,9 @@ async def custom(
         if not isinstance(mode, list):
             mode = [mode]
 
+        if format_amba is True:
+            value = self._fmt_amba(address, size, value)
+
         # Need to copy data as we'll have to shift address/size
         t_address = copy.deepcopy(address)
         t_value = copy.deepcopy(value)
diff --git a/cocotbext/ahb/version.py b/cocotbext/ahb/version.py
index a34b2f6..a3a9bd5 100644
--- a/cocotbext/ahb/version.py
+++ b/cocotbext/ahb/version.py
@@ -1 +1 @@
-__version__ = "0.4.7"
+__version__ = "0.4.8"
diff --git a/tests/test_ahb.py b/tests/test_ahb.py
index 6442363..5fadcd0 100644
--- a/tests/test_ahb.py
+++ b/tests/test_ahb.py
@@ -4,7 +4,7 @@
 # License           : MIT license <Check LICENSE>
 # Author            : Anderson I. da Silva (aignacio) <anderson@aignacio.com>
 # Date              : 08.10.2023
-# Last Modified Date: 09.09.2024
+# Last Modified Date: 27.12.2024
 
 import cocotb
 import os
@@ -86,7 +86,7 @@ async def run_test(dut, bp_fn=None, pip_mode=False):
     txn_type = [pick_random_value([1, 0]) for _ in range(N)]
 
     resp = await ahb_master.custom(
-        address, value, txn_type, size, pip_mode
+        address, value, txn_type, size, pip_mode, format_amba=True
     )