@ptrblck I have my code below. You would just need the CIFAR-10 dataset . To run this I would have to send a .pth file but unable to send the file herein. How do I send it?
model=torch.load('/content/cifar_fullprecison_vgg19_shayer_50.pth')
class Ternary_batch_rel(torch.nn.Module):
  def __init__(self,batchnorm_size):
    super(Ternary_batch_rel,self).__init__()
    self.l1=torch.nn.Sequential(
    torch.nn.ReLU(),
    torch.nn.BatchNorm2d(batchnorm_size)
    )
  
  def forward(self,x):
    out=self.l1(x)
    return out
    
z1=Ternary_batch_rel(128).to(device)
z2=Ternary_batch_rel(256).to(device)
z3=Ternary_batch_rel(512).to(device)
​
class Ternary_max_pool(torch.nn.Module):
  def __init__(self):
    
    super(Ternary_max_pool,self).__init__()
    self.l1=torch.nn.Sequential(
    torch.nn.MaxPool2d(kernel_size=2,stride=2))
      
  def forward(self,x):
    out=self.l1(x)
    return out
​
zm=Ternary_max_pool().to(device)
sigm=torch.nn.Sigmoid().to(device)
def convlayer1param():
  s1=torch.zeros([128,3,3,3])
  s1=model['layer1.0.weight']
  return s1
 
def convlayer2param():
  s1=torch.zeros([128,128,3,3])
  s1=model['layer1.3.weight']
  return s1

def convlayer3param():
  s1=torch.zeros([256,128,3,3])
  s1=model['layer2.0.weight']
  return s1
def convlayer4param():
  s1=torch.zeros([256,256,3,3])
  s1=model['layer2.3.weight']
  return s1

def convlayer5param():
  s1=torch.zeros([512,256,3,3])
  s1=model['layer3.0.weight']
  return s1

def convlayer6param():
  s1=torch.zeros([512,512,3,3])
  s1=model['layer3.3.weight']
  return s1
def reparamcnn1(a,b,h):
  weight_m= (2*sigm(b)-(2*sigm(a)*sigm(b))-1+sigm(a))
  weight_v=(1-sigm(a))-weight_m**2
  om=F.conv2d(h,weight_m,padding=1)
  ov=F.conv2d(h**2,weight_v,padding=1)
  e=torch.randn(ov.shape).cuda()
  z=om+ov*e
  return z1(z)

def reparamcnn2(a,b,h):
  weight_m=(2*sigm(b)-(2*sigm(a)*sigm(b))-1+sigm(a))
  weight_v=(1-sigm(a))-weight_m**2
  om=F.conv2d(h,weight_m,padding=1)
  ov=F.conv2d(h**2,weight_v,padding=1)
  e=torch.randn(ov.shape).cuda()
  #e=torch.randn(1).cuda()
  z=om+ov*e
  op=z1(z)
  return zm(op)

def reparamcnn3(a,b,h):
  weight_m= (2*sigm(b)-(2*sigm(a)*sigm(b))-1+sigm(a))
  weight_v=(1-sigm(a))-weight_m**2
  om=F.conv2d(h,weight_m,padding=1)
  ov=F.conv2d(h**2,weight_v,padding=1)
  e=torch.randn(ov.shape).cuda()
  #e=torch.randn(1).cuda()
  z=om+ov*e
  return z2(z)
  
def reparamcnn4(a,b,h):
  weight_m=(2*sigm(b)-(2*sigm(a)*sigm(b))-1+sigm(a))
  weight_v=(1-sigm(a))-weight_m**2
  om=F.conv2d(h,weight_m,padding=1)
  ov=F.conv2d(h**2,weight_v,padding=1)
  e=torch.randn(ov.shape).cuda()
  #e=torch.randn(1).cuda()
  z=om+ov*e
  op=z2(z)
  return zm(op)
 
def reparamcnn5(a,b,h):
  weight_m= (2*sigm(b)-(2*sigm(a)*sigm(b))-1+sigm(a))
  weight_v=(1-sigm(a))-weight_m**2
  
  om=F.conv2d(h,weight_m,padding=1)
  ov=F.conv2d(h**2,weight_v,padding=1)
  e=torch.randn(ov.shape).cuda()
  #e=torch.randn(1).cuda()
  z=om+ov*e

  return z3(z)

def reparamcnn6(a,b,h):
  weight_m=(2*sigm(b)-(2*sigm(a)*sigm(b))-1+sigm(a))
  weight_v=(1-sigm(a))-weight_m**2
  om=F.conv2d(h,weight_m,padding=1)
  ov=F.conv2d(h**2,weight_v,padding=1)
  e=torch.randn(ov.shape).cuda()
  #e=torch.randn(1).cuda()
  z=om+ov*e
  op=z3(z)
  return zm(op)

def reparamfc(a,b,h):
  weight_m=(2*sigm(b)-(2*sigm(a)*sigm(b))-1+sigm(a))
  weight_v=(1-sigm(a))-weight_m**2
  om=torch.matmul(weight_m,h)
  ov=torch.matmul(weight_v,h**2)
  e=torch.randn(ov.shape).cuda()
 # e=torch.randn(1).cuda()
  z=om+ov*e
  return z

def initialize(wfp):
  wtilde=wfp/torch.std(wfp)
  sigma_a=0.95-((0.95-0.05)*torch.abs(wtilde))
  sigma_b=0.5*(1+(wfp/(1-sigma_a)))
  sigma_a=torch.clamp(sigma_a,0.05,0.95)
  sigma_b=torch.clamp(sigma_b,0.05,0.95)
  a=torch.log(sigma_a/(1-sigma_a)).requires_grad_().cuda()
  b=torch.log(sigma_b/(1-sigma_b)).requires_grad_().cuda()
  
  return a,b

w1fpconv=convlayer1param()
w2fpconv=convlayer2param()
w3fpconv=convlayer3param()
w4fpconv=convlayer4param()
w5fpconv=convlayer5param()
w6fpconv=convlayer6param()
wfp1=model['layer4.1.weight']
wfp2=model['layer4.4.weight']
al1,bl1=initialize(w1fpconv)
al2,bl2=initialize(w2fpconv)
al3,bl3=initialize(w3fpconv)
al4,bl4=initialize(w4fpconv)
al5,bl5=initialize(w5fpconv)
al6,bl6=initialize(w6fpconv)
a1,b1=initialize(wfp1)
a2,b2=initialize(wfp2)

al1=torch.nn.Parameter(al1)
bl1=torch.nn.Parameter(bl1)
al2=torch.nn.Parameter(al2)
bl2=torch.nn.Parameter(bl2)
al3=torch.nn.Parameter(al3)
bl3=torch.nn.Parameter(bl3)
al4=torch.nn.Parameter(al4)
bl4=torch.nn.Parameter(bl4)
al5=torch.nn.Parameter(al5)
bl5=torch.nn.Parameter(bl5)
al6=torch.nn.Parameter(al6)
bl6=torch.nn.Parameter(bl6)
a1=torch.nn.Parameter(a1)
b1=torch.nn.Parameter(b1)
a2=torch.nn.Parameter(a2)
b2=torch.nn.Parameter(b2)


betaparam=1e-11
lossfunc=torch.nn.CrossEntropyLoss().to(device)

lr=0.01
optimizer=torch.optim.Adam([al1,bl1,al2,bl2,al3,bl3,al4,bl4,al5,bl5,al6,bl6,a1,b1,a2,b2],lr,weight_decay=1e-4)

num_epochs=10

for epoch in range(num_epochs):
  for i,(images,labels) in enumerate(train_loader):
    images=images.to(device)
    labels=labels.to(device)
    y1=reparamcnn1(al1,bl1,images)
    y2=reparamcnn2(al2,bl2,y1)
    y3=reparamcnn3(al3,bl3,y2)
    y4=reparamcnn4(al4,bl4,y3)
    y5=reparamcnn5(al5,bl5,y4)
    y6=reparamcnn6(al6,bl6,y5)
    y6=y6.reshape(y6.size(0),-1)
    y6=torch.t(y6)
    y7=F.dropout(y6)
    y8=reparamfc(a1,b1,y7)
    y9=F.relu(y8)
    y10=F.dropout(y9)
    yout=reparamfc(a2,b2,y10)
    yout=torch.t(yout)
    #yout=F.softmax(yout,dim=1)
    l2=al1.norm(2)+bl1.norm(2)+al2.norm(2)+bl2.norm(2)+al3.norm(2)+bl3.norm(2)+al4.norm(2)+bl4.norm(2)+al5.norm(2)+bl5.norm(2)+al6.norm(2)+bl6.norm(2)+a1.norm(2)+b1.norm(2)+a2.norm(2)+b2.norm(2)
    lossi=lossfunc(yout,labels)+(betaparam*l2)
    if(epoch==170):
      lr=0.001
      for param_group in optimizer.param_groups:
        param_group['lr']=lr  
    lossi.backward()
    optimizer.step()
    optimizer.zero_grad()
  print('epoch {}'.format(epoch),'loss = {}'.format(lossi.item()))