np.argwhere -> np.argmax

[AndreiMoraru123]

What does this PR do?

Solves: #535

Related Issue(s):

Type of Change

• Bug fix (non-breaking change that fixes an issue)

Testing

• I have tested this change locally
• I have added/updated tests for this change

Test details:

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code where necessary, particularly in hard-to-understand areas
• My changes generate no new warnings or errors
• I have updated the documentation accordingly (if applicable)

Additional Context

[CLAassistant]

All committers have signed the CLA.