Merge branch 'main' into quantization

bayesian_torch/layers/flipout_layers/conv_flipout.py

@@ -40,6 +40,7 @@
 from torch.quantization.observer import HistogramObserver, PerChannelMinMaxObserver, MinMaxObserver
 from torch.quantization.qconfig import QConfig
 
+
 from torch.distributions.normal import Normal
 from torch.distributions.uniform import Uniform
 
@@ -419,6 +420,7 @@ def forward(self, x, return_kl=True):
         return out
 
 
+
 class Conv3dFlipout(BaseVariationalLayer_):
     def __init__(self,
                  in_channels,

bayesian_torch/layers/flipout_layers/linear_flipout.py

@@ -195,3 +195,4 @@ def forward(self, x, return_kl=True):
         if return_kl:
             return out, kl
         return out
+

bayesian_torch/layers/variational_layers/conv_variational.py

@@ -374,7 +374,7 @@ def forward(self, input, return_kl=True):
             else:
                 kl = kl_weight
             return out, kl
-
+            
         return out
 
 
@@ -973,12 +973,3 @@ def forward(self, input, return_kl=True):
 
         return out
 
-if __name__=="__main__":
-    m = Conv2dReparameterization(3,3,3)
-    m.eval()
-    m.prepare()
-    m.qconfig = torch.quantization.get_default_qconfig("fbgemm")
-    mp = torch.quantization.prepare(m)
-    input = torch.randn(3,3,4,4)
-    mp(input)
-    mq = torch.quantization.convert(mp)

bayesian_torch/layers/variational_layers/linear_variational.py

@@ -162,6 +162,7 @@ def forward(self, input, return_kl=True):
         tmp_result = sigma_weight * eps_weight
         weight = self.mu_weight + tmp_result
 
+
         if return_kl:
             kl_weight = self.kl_div(self.mu_weight, sigma_weight,
                                     self.prior_weight_mu, self.prior_weight_sigma)