I have a model wherein I need to compute the cross-entropy loss for only a certain index. The node_mask allows for the selection of those indices. There is also ignore_index field in the Cross-Entropy loss class. Since I have the gt for only these masked labels, I thought maybe it would make sense in terms of the overall space efficiency if we just do the masked selection beforehand.
x = self.linear1(self.node_features)  # W, 200
x = self.linear2(x)  # W, 1024
x = F.relu(x)
x = x[self.node_mask] # W, 14

However, I am still not sure if this is going to lead to the same effect or I should stick to ignore_index term and expand my Ground truth to a 1024 dimensional vector with -100 at the non-mask indices.
Any help would be highly appreciated.