Fix KL computation

[y0ast]

Hello,

I think there might be a few problems in your model definitions.

In particular:

• in resnet_flipout.py and resnet_variational.py you only sum the kl of the last block inside self.layerN
• in resnet_flipout_large.py and resnet_variational_large.py you check for is None while you probably want is not None or actually no check at all since it can't be None in any reasonable setting. Also the str(layer) check is odd since it contains a BasicBlock or BottleNeck object (you're looping over an nn.Sequential of blocks). In fact in this code that string check is very likely superfluous (didn't test this, but I did include it in this PR as example)

I hope you can confirm and perhaps fix these issues, which will help me (and maybe others) in building on your nice codebase :)

[ranganathkrishnan]

Thank you for the fixes! The kl_sum addition bug was in the list to be fixed. The str(layer) check was present as I was using BatchNorm2d layer directly from torch.nn in the sequential block. Now that check is not necessary as the wrapper from bayesian_torch.layers.batchnorm is used.

[ranganathkrishnan]

Hello @y0ast ,

Thank you for your contribution to the repo. These bugs were in the list to be fixed along with other fixes. Thanks for the PR, merging the commit.

[y0ast]

Thanks for merging! Are you able to share the list of other fixes you are planning? (So I can avoid those bugs in my own work building on your codebase).

[ranganathkrishnan]

I have pushed the commits (09b8d0a and 6a370d9) to the repo just before merging your PR. If you checkout from the head of main branch, these fixes should be already in.