#这里的self.layers就是论文中N层MHA连接FF的结构
self.layers = nn.Sequential(*(
MultiHeadAttentionLayer(n_heads, embed_dim, feed_forward_hidden, normalization)
for _ in range(n_layers)
# Batch multiply to get initial embeddings of nodes
#将初始的self.init_embed通过刚刚定义的layers变换到h和h`
h = self.init_embed(x.view(-1, x.size(-1))).view(*x.size()[:2], -1) if self.init_embed is not None else x
h = self.layers(h) 返回的是`h`和`h.meaan`，也就是上面说的，node embeddings和graph embedding。

SkipConnection(
    nn.Sequential(
        nn.Linear(embed_dim, feed_forward_hidden),
        nn.ReLU(),
        nn.Linear(feed_forward_hidden, embed_dim)
    ) if feed_forward_hidden > 0 else nn.Linear(embed_dim, embed_dim)
),
Normalization(embed_dim, normalization)  

    batch_size = state.ids.size(0)

    # Perform decoding steps
    i = 0
    while not (self.shrink_size is None and state.all_finished()):

        if self.shrink_size is not None:
            unfinished = torch.nonzero(state.get_finished() == 0)
            if len(unfinished) == 0:
                break
            unfinished = unfinished[:, 0]
            # Check if we can shrink by at least shrink_size and if this leaves at least 16
            # (otherwise batch norm will not work well and it is inefficient anyway)
            if 16 <= len(unfinished) <= state.ids.size(0) - self.shrink_size:
                # Filter states
                state = state[unfinished]
                fixed = fixed[unfinished]
		# **根据fixed和state得到mask**
        log_p, mask = self._get_log_p(fixed, state) 最关键的，mask虽然在`self._get_log_p()`这里返回的，但在这个方法中，`mask = state.get_mask()`得到了mask，所以还是得去`state_cvrp.py`中去看，也就是：  
def get_mask(self):
    """
    Gets a (batch_size, n_loc + 1) mask with the feasible actions (0 = depot), depends on already visited and
    remaining capacity. 0 = feasible, 1 = infeasible
    Forbids to visit depot twice in a row, unless all nodes have been visited
    :return:
    """

    if self.visited_.dtype == torch.uint8:
        visited_loc = self.visited_[:, :, 1:]
    else:
        visited_loc = mask_long2bool(self.visited_, n=self.demand.size(-1))

    # Nodes that cannot be visited are already visited or too much demand to be served now
    mask_loc = (
        visited_loc |
        # For demand steps_dim is inserted by indexing with id, for used_capacity insert node dim for broadcasting
        (self.demand[self.ids, :] + self.used_capacity[:, :, None] > self.VEHICLE_CAPACITY)
    )

    # Cannot visit the depot if just visited and still unserved nodes
    mask_depot = (self.prev_a == 0) & ((mask_loc == 0).int().sum(-1) > 0)
    return torch.cat((mask_depot[:, :, None], mask_loc), -1) 我理解的是里面应该用了位操作的方式，具体没有细看（因为太难了哭泣。。。），欢迎各位大佬交流。反正就是0是没有mask，1是mask。后面的就不难懂了：  
        # Select the indices of the next nodes in the sequences, result (batch_size) long
        selected = self._select_node(log_p.exp()[:, 0, :], mask[:, 0, :])  # Squeeze out steps dimension
		#选好的结点进行更新，以避免下次再选
        state = state.update(selected)

        # Now make log_p, selected desired output size by 'unshrinking'
        if self.shrink_size is not None and state.ids.size(0) < batch_size:
            log_p_, selected_ = log_p, selected
            log_p = log_p_.new_zeros(batch_size, *log_p_.size()[1:])
            selected = selected_.new_zeros(batch_size)

            log_p[state.ids[:, 0]] = log_p_
            selected[state.ids[:, 0]] = selected_

        # Collect output of step
		#输出的概率和pi，也就是序列
        outputs.append(log_p[:, 0, :])
        sequences.append(selected)

        i += 1

    # Collected lists, return Tensor
    return torch.stack(outputs, 1), torch.stack(sequences, 1)  

training_dataset = baseline.wrap_dataset(problem.make_dataset(
    size=opts.graph_size, num_samples=opts.epoch_size, distribution=opts.data_distribution))
training_dataloader = DataLoader(training_dataset, batch_size=opts.batch_size, num_workers=1)  

# Evaluate baseline, get baseline loss if any (only for critic)
bl_val, bl_loss = baseline.eval(x, cost) if bl_val is None else (bl_val, 0)

# Calculate loss
reinforce_loss = ((cost - bl_val) * log_likelihood).mean()
loss = reinforce_loss + bl_loss

# Perform backward pass and optimization step
optimizer.zero_grad()
loss.backward()
# Clip gradient norms and get (clipped) gradient norms for logging
grad_norms = clip_grad_norms(optimizer.param_groups, opts.max_grad_norm)
optimizer.step()