anantguptadbl:

self.weight = nn.Parameter(torch.randint(high = 2,size=(128,128)).float())


In that case you can replace
self.weight = nn.Parameter(torch.randint(high = 2,size=(128,128)).float())
with
self.weight = nn.Parameter(torch.randint(high = 2,size=(M,128,128)).float())