Batch data selection at intermediate layer

autograd
1

Dear community,

I want to implement the model where a certain batch selection occurs at the intermediate layer of the model. The selection outputs the subset of the current batch data and feed it to the classifier. Below is the model code

class ConvNet(Module):
    def __init__(self, input_shape = (28,28), num_classes=62):
        self.featdim = 160
        self.num_classes = num_classes
        self.input_shape = input_shape
        
        super(ConvNet, self).__init__()
        self.conv1 = Conv2d(1, 5, 5)
        self.pool = MaxPool2d(2, 2)
        self.conv2 = Conv2d(5, 10, 5)
        self.fc1 = Linear(self.featdim, 129)
        self.fc2 = Linear(129, self.num_classes)

    def get_feature(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(x.shape[0],-1)
        return x

    def classifier(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
        
    def forward(self, x, repr=True):
        if repr:
            x = self.get_feature(x)
        x = self.classifier(x)
        return x

and the main code

m = ConvNet()
d = datasets.MNIST(root='dataset', train=True, transform=transforms.ToTensor())
l = DataLoader(d, 8)
for X, y in l:
    h = m.get_feature(X)
    idx_sel = torch.randperm(8)[:4]
    y_s = y[idx_sel]
    h_s = h[idx_sel]
    c = m.classifier(h_s.view(4, -1))
    loss = torch.nn.CrossEntropyLoss()(c, y_s)
    loss.backward()

I am unsure whether the unselected data will also affect the feature layer during the backpropagation. Is this essentially the same as training data directly on X[idx_sel]? Thanks.