diff --git a/addons/mcs_payroll_localization/hooks.py b/addons/mcs_payroll_localization/hooks.py
index e153c58..c32ff67 100644
--- a/addons/mcs_payroll_localization/hooks.py
+++ b/addons/mcs_payroll_localization/hooks.py
@@ -21,9 +21,9 @@ class PayrollSetup:
 
     def run(self):
         categories = self._ensure_categories()
-        for company in self.Company.search([("country_code", "=", "IN")]):
+        for company in self.Company.search([("country_id.code", "=", "IN")]):
             self._ensure_india(company, categories)
-        for company in self.Company.search([("country_code", "=", "CA")]):
+        for company in self.Company.search([("country_id.code", "=", "CA")]):
             self._ensure_canada(company, categories)
 
     def _ensure_categories(self):