Thanks for your suggestion. I passed column index to the forward method. It worked well and sped up a lot.
    def forward(self, x, col_idx):
        
        for k, layer in enumerate(self.layers):
            x = layer(x)
            if self.act_function[k] == 'sigmoid':
                x = torch.sigmoid(x)
            elif self.act_function[k] == 'relu':
                x = F.relu(x)
            elif self.act_function[k] == 'leaky_relu':
                x = F.leaky_relu(x)
            
            if self.batch_norm == True:
                x = self.bns[k](x)
                
            x = self.dropouts[k](x)
        
        outputs = torch.zeros([fps.shape[0], len(col_idx)])
        
        for k, idx in enumerate(col_idx):
            outputs[:, k] = torch.squeeze(self.heads[idx](x), 1)
        """
        outputs = []
        for i, head in enumerate(self.heads):
            if i in col_idx:
                outputs.append(head(x))
        """
        #outputs = torch.cat(outputs).view(len(outputs), -1).T
        return outputs