Classifier using pytorch
Choose lr of optimizer something very small. It might be because of exploding gradient.

In self.weight use nn.Parameter() then pass your torch.zeros() to make it a model parameter.