As you say, the output of a $Q$ network is typically a value for all actions of the given state. Let us call this output $\mathbf{x} \in \mathbb{R}^{|\mathcal{A}|}$. To train your network using the squared bellman error you need first calculate the scalar target $y = r(s, a) + \max_a Q(s', a)$. Then, to train the network we take a vector $\mathbf{x'} = \mathbf{x}$ and change the $a$th element of it to be equal to $y$, where $a$ is the action you took in state $s$; call this modified vector $\mathbf{x'}_a$. We calculate the loss $\mathcal{L}(\mathbf{x}, \mathbf{x'}_a)$ and back propagate through this to update the parameters of our network.
Note that when we use $Q$ to calculate $y$ we typically use some form of target network; this can be a copy of $Q$ where the parameters are only updated every $i$th update or a network whose weights are updated using a polyak average with the main networks weights after every update.
Judging by your code it looks as though your action selection is what might be causing you some problems. As far as I can tell you're always acting greedily with respect to your $Q$-function. You should be looking to act $\epsilon$-greedily, i.e. with probability $\epsilon$ take a random action and act greedily otherwise. Typically you start with $\epsilon=1$ and decay it each time a random action is taken down to some small value such as 0.05.