Yes, that's a classical proof. It can be recast into the following proof comparing parity of powers of $2$ in factorizations (a special case of uniqueness of prime factorizations). Suppose $\, A^{\large 2} = 2B^{\large 2}$
It is straightforward to prove that every natural can be written uniquely in the form $\, 2^{\large k} n\,$ for odd $\,n.\,$ So we have $\, A= 2^{\large J} a,\ B = 2^{\large K} b\,$ for $\,a,b\,$ odd. Therefore $\, A^{\large 2}\! = 2 B^{\large 2}$ $\Rightarrow$ $\,2^{\color{#c00}{\large 2J}} a^{\large 2} = 2^{\color{#0a0}{\large 1+2K}} b^{\large 2}.\ $ But $\,a,b\,$ odd $\,\color{#f2f}\Rightarrow$ $\,a^{\large 2},b^{\large 2}\,$ odd too, which contradicts uniqueness (LHS has $\rm\color{#c00}{even}$ vs. RHS $\rm\color{#0a0}{odd}$ power of $2).$
To generalize this proof from $\,p = 2\,$ to an arbitrary prime $\,p\,$ would require generalizing the $\,\rm odd^{\large 2} = odd\,$ step that we used above, $ $ i.e. we would need a proof of the following inference $\ p\nmid a,b\,\color{#f2f}\Rightarrow\, p\nmid a^{\large 2},b^{\large 2}.\ $ Just as for the case $\,p=2,\,$ for any fixed prime $\,p\,$ we could prove it by a brute-force case analysis. But to prove the result for all primes $\,p\,$ we need something much deeper, e.g. Euclid's lemma $\ p\mid ab\,\Rightarrow\, p\mid a\,$ or $\,p\mid b\ $ (or closely related results).