Thanks for the reply, I was not aware of that. However, even when I check the .grad during forward, after the first time, it appears that most of my parameters have grad=torch.zeros(some_size), and the parameters which belong to the above method have grad=None. These parameters do not change during training, unlike the other parameters of my model.
I’m using pytorch lightning, so I implemented a forward function, attached here:
    def forward(self, image, reads_metadata, metadata, maternal_labels,
                paternal_labels):
        cfDNA_reads, parental_reads = torch.split(image, 3, dim=1)

        attention_output = self.soft_self_attention_layer(
            K=cfDNA_reads,
            Q=torch.cat([cfDNA_reads, parental_reads], dim=2)
        )

        cnn_input = torch.cat([attention_output, parental_reads], dim=1)

        x1 = self.cnn(cnn_input)

        x2 = metadata

        bayesian_logits = self.get_bayesian_predictions(
            maternal_labels=maternal_labels,
            paternal_labels=paternal_labels,
            reads_metadata=reads_metadata,
            cffdna_ratio_mean=\
            metadata[:, POS_META_FEATURES.index(CFFDNA_RATIO_MEAN)]
        )

        x = torch.cat((x1, x2, bayesian_logits.cuda()), dim=1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x, bayesian_logits