Project: 014_Energy-based_Backdoor_Defense_Against_Federated_Gr | Function: adjust_bn_layers (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Meta Energy Calibration method that optimizes batch normalization parameters to enhance energy-based defense against backdoor attacks

**KEY TERMS:** {
  "p_data": "Me(vi) energy of benign samples",
  "p_neigh": "M̃e(ṽi^adv) energy of perturbed samples",
  "energy": "Ms_θ(vi) meta energy score",
  "energy_grad": "∇vi Ms_θ(vi) energy gradient",
  "energy_grad_inner": "∇vi Ms_θ(vi)^T ∇vi Ms_θ(vi) gradient inner product",
  "energy_squared_sum": "||Ms_θ(vi)||^2 energy magnitude term"
}

**CONTEXT CODE:**             optimizer.zero_grad()
            output = self.forward(self.features, self.edge_index, self.edge_weight)
            loss_train = F.nll_loss(output[idx_train], labels[idx_train])
            loss_train.backward()
            optimizer.step()
            if verbose and i % 10 == 0:
                print('Epoch {}, training loss: {}'.format(i, loss_train.item()))

        self.eval()
        output = self.forward(self.features, self.edge_index, self.edge_weight)
        self.output = output

    def _train_with_val(self, global_model, features, labels, idx_train, idx_val, edge_index, edge_weight,
                        aug_edge_index, aug_edge_weight, train_iters, args, verbose):
        if verbose:
            print('=== training gcn model ===')
        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
        best_loss_val = 100
        best_acc_val = -10

        for i in range(train_iters):
            set_random_seed(42 + i)  # Fixed seed for testing
            self.train()
            optimizer.zero_grad()
            output = self.forward(features, edge_index, edge_weight)
            loss_train = F.nll_loss(output[idx_train], labels[idx_train])

            total_loss = loss_train
            total_loss.backward()
            optimizer.step()

            self.eval()
            with torch.no_grad():
                output = self.forward(features, edge_index, edge_weight)
                loss_val = F.nll_loss(output[idx_val], labels[idx_val])
                acc_val = accuracy(output[idx_val], labels[idx_val])
                acc_train = accuracy(output[idx_train], labels[idx_train])

            if acc_val > best_acc_val:
                best_acc_val = acc_val
                self.output = output
                weights = deepcopy(self.state_dict())

        if verbose:
            print('=== picking the best model according to the performance on validation ===')
        self.load_state_dict(weights)

        return loss_train.item(), loss_val.item(), acc_train, acc_val


    def adjust_bn_layers(self, features, edge_index, edge_weight, aug_edge_index, aug_edge_weight):
        raise NotImplementedError("Function implementation has been masked for testing")
    def test(self, features, edge_index, edge_weight, labels, idx_test):
        self.eval()
        with torch.no_grad():
            output = self.forward(features, edge_index, edge_weight)
            acc_test = accuracy(output[idx_test], labels[idx_test])

        return float(acc_test)

    def test_with_correct_nodes(self, features, edge_index, edge_weight, labels, idx_test):
        self.eval()
        output = self.forward(features, edge_index, edge_weight)
        correct_nids = (output.argmax(dim=1)[idx_test] == labels[idx_test]).nonzero().flatten()  # return a tensor
        acc_test = accuracy(output[idx_test], labels[idx_test])
        return acc_test, correct_nids


Complete the function implementation replacing "raise NotImplementedError".