diff --git a/rma/hooks.py b/rma/hooks.py
index 9aafcaf0..ed625906 100644
--- a/rma/hooks.py
+++ b/rma/hooks.py
@@ -56,15 +56,10 @@ def post_init_hook(cr, registry):
                 )
                 data[picking_type] = stock_picking_type.create(values).id
 
-        rma_out_type = stock_picking_type.browse(data["rma_out_type_id"])
-        rma_out_type.write(
-            {"return_picking_type_id": data.get("rma_in_type_id", False)}
-        )
-        rma_in_type = stock_picking_type.browse(data["rma_in_type_id"])
-        rma_in_type.write(
-            {"return_picking_type_id": data.get("rma_out_type_id", False)}
-        )
-        whs.write(data)
+        if data:
+            whs.write(data)
+        whs.rma_in_type_id.return_picking_type_id = whs.rma_out_type_id.id
+        whs.rma_out_type_id.return_picking_type_id = whs.rma_in_type_id.id
 
     # Create rma locations and picking types
     warehouses = env["stock.warehouse"].search([])