diff --git a/addons/stock/data/stock_data.yml b/addons/stock/data/stock_data.yml
index 30289281762..6ab264bb78f 100644
--- a/addons/stock/data/stock_data.yml
+++ b/addons/stock/data/stock_data.yml
@@ -1,6 +1,3 @@
--
-  !python {model: res.partner, id: base.main_partner}: |
-    main_warehouse = self.env['stock.warehouse'].browse(ref('warehouse0'))
 -
   !python {model: ir.model.data, id: False}: |
     main_warehouse = self.env['stock.warehouse'].browse(ref('warehouse0'))
@@ -28,7 +25,3 @@
     companies = self.search([('internal_transit_location_id', '=', False)])
     for company in companies:
         company.create_transit_location()
-        internal_transit = company.internal_transit_location_id.id
-        warehouses = self.env['stock.warehouse'].search([('partner_id', '=', company.partner_id.id)])
-        warehouses.mapped('partner_id').with_context(force_company=company.id).write({'property_stock_customer': internal_transit, 'property_stock_supplier': internal_transit})
-
diff --git a/addons/stock/models/res_company.py b/addons/stock/models/res_company.py
index 94a7b58b711..0fbf2868725 100644
--- a/addons/stock/models/res_company.py
+++ b/addons/stock/models/res_company.py
@@ -28,6 +28,12 @@ class Company(models.Model):
         location.sudo().write({'company_id': self.id})
         self.write({'internal_transit_location_id': location.id})
 
+        warehouses = self.env['stock.warehouse'].search([('partner_id', '=', self.partner_id.id)])
+        warehouses.mapped('partner_id').with_context(force_company=self.id).write({
+            'property_stock_customer': location.id,
+            'property_stock_supplier': location.id,
+        })
+
     @api.model
     def create(self, vals):
         company = super(Company, self).create(vals)