def METHOD_NAME(self, data_train, data_val):
    input_train, target_train = data_train
    input_val, target_val = data_val
    # sample weights (alphas) from the dirichlet distribution (parameterized by beta) and set to edges
    self.graph.update_edges(
        update_func=lambda edge: self.sample_alphas(edge),
        scope=self.scope,
        private_edge_data=False,
    )
    # Update architecture weights
    self.arch_optimizer.zero_grad()
    logits_val = self.graph(input_val)
    val_loss = self.loss(logits_val, target_val)
    if self.reg_type == "kl":
        val_loss += self._get_kl_reg()
    val_loss.backward()
    if self.grad_clip:
        torch.nn.utils.clip_grad_norm_(
            self.architectural_weights.parameters(), self.grad_clip
        )
    self.arch_optimizer.METHOD_NAME()
    # has to be done again, cause val_loss.backward() frees the gradient from sampled alphas
    # TODO: this is not how it is intended because the samples are now different. Another
    # option would be to set val_loss.backward(retain_graph=True) but that requires more memory.
    self.graph.update_edges(
        update_func=lambda edge: self.sample_alphas(edge),
        scope=self.scope,
        private_edge_data=False,
    )
    # Update op weights
    self.op_optimizer.zero_grad()
    logits_train = self.graph(input_train)
    train_loss = self.loss(logits_train, target_train)
    train_loss.backward()
    if self.grad_clip:
        torch.nn.utils.clip_grad_norm_(self.graph.parameters(), self.grad_clip)
    self.op_optimizer.METHOD_NAME()
    # in order to properly unparse remove the alphas again
    self.graph.update_edges(
        update_func=self.remove_sampled_alphas,
        scope=self.scope,
        private_edge_data=False,
    )
    return logits_train, logits_val, train_loss, val_loss