diff --git a/account_payment_partner_default_bank/models/account_invoice.py b/account_payment_partner_default_bank/models/account_invoice.py
index eec45d5..0145038 100644
--- a/account_payment_partner_default_bank/models/account_invoice.py
+++ b/account_payment_partner_default_bank/models/account_invoice.py
@@ -18,8 +18,9 @@ class AccountInvoice(models.Model):
             company = self.env.user.company_id
         res = super(AccountInvoice, self)._onchange_partner_id()
         if self.type == 'out_invoice' and not self.partner_bank_id:
-            if company.bank_ids:
-                self.partner_bank_id = company.bank_ids[0]
+            bank_id = self._get_partner_bank_id(company.id)
+            if bank_id:
+                self.partner_bank_id = bank_id
 
         return res
 
@@ -31,6 +32,7 @@ class AccountInvoice(models.Model):
             company = self.env.user.company_id
         res = super(AccountInvoice, self)._onchange_payment_mode_id()
         if self.type == 'out_invoice' and not self.partner_bank_id:
-            if company.bank_ids:
-                self.partner_bank_id = company.bank_ids[0]
+            bank_id = self._get_partner_bank_id(company.id)
+            if bank_id:
+                self.partner_bank_id = bank_id
         return res