Here is a slight modification of the code
import torch
from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
import torchvision.models as models
import torch.nn.functional as F

class BinaryFilterLayer(nn.Module):
  def __init__(self,number_of_features):
    super(BinaryFilterLayer, self).__init__()
    self.number_of_features = number_of_features
    # Generating the initial binary filters (only 0 and 1)
    self.weight = nn.Parameter(torch.randint(high = 2,size=(128,128)).float())
    self.sigm=nn.Sigmoid()

  def forward(self, input):
      #print("size of input is {0} and size of weight is {1}".format(input.size(),self.weight.size()))
      self.weight
      output = input * torch.round(self.sigm(self.weight))
      return output

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(128,128)


class Model(nn.Module):
  def __init__(self,encoding_dim):
    super(Model, self).__init__()
    self.encoding_dim = encoding_dim
    # Encode
    self.binary_filters = BinaryFilterLayer(self.encoding_dim)
    # Decode
    self.decoder = nn.Sequential(
                    #nn.LayerNorm(self.encoding_dim),
                    #nn.Linear(self.encoding_dim,128*128),
                    #nn.LayerNorm(128*128),
                    #nn.ReLU(),
                    #Flatten(),
                    nn.Conv2d(1,64,kernel_size= 9,padding=2),
                    nn.BatchNorm2d(64),
                    nn.ReLU(),
                    nn.Conv2d(64,32,kernel_size= 1,padding=2),
                    nn.BatchNorm2d(32),
                    nn.ReLU(),
                    nn.Conv2d(32,1,kernel_size= 5,padding=2),
                    nn.BatchNorm2d(1),
                    nn.Sigmoid()
                    )

  def sumfunc(self,x):
    x = torch.sum(x,dim=3)     # maybe there is an efficient way but this should work well also i guess
    x = torch.sum(x,dim=2)
    return x/(128*128)            # not sure if i need to normalize here but probably not
  
  def forward(self, x):
    x = self.binary_filters(x)  # multiply with binary layers
    #print("Post binary filters, the size of x is {0}".format(x.size()))
    #x = self.sumfunc(x)        # get the sum of each layer
    x = self.decoder(x.view(1,1,128,128))        # decode
    #print("After decoder the size of x is {0}".format(x.size()))
    return x 

model=Model(125)
criterion=torch.nn.MSELoss()
optimizer=torch.optim.Adam(model.parameters(),lr=0.001)
epochs=1000
from skimage import io
import cv2
labrador = io.imread("gdrive/My Drive/Colab Notebooks/YellowLabradorLooking_new.jpg", as_gray=True)
labrador=cv2.resize(labrador,(128,128))
labrador=Variable(torch.from_numpy(labrador.astype(np.float32)))
print("Input size is {0}".format(labrador.size()))

for curEpoch in range(epochs):
  model.zero_grad()
  output=model(labrador)
  loss=criterion(labrador,output)
  loss.backward()
  optimizer.step()
  if(curEpoch %100==0):
    print("Loss is {0}".format(loss.item()))

import matplotlib.pyplot as plt
%matplotlib inline
plt.imshow(labrador.detach().numpy().reshape(128,128))
plt.show()
plt.imshow(model.binary_filters.weight.detach().numpy() * labrador.detach().numpy().reshape(128,128))



Hope this helps