kornia 0.4.0 - torch 1.6.0 releases

[KeremTurgutlu]

I tried the sample code available in README, it looks like library depends on an earlier version of kornia 0.3.2. Since kornia just very recently released 0.4.0 it might be a bit of an early issue but I figured it would be helpful. Also new version of kornia 0.4.0 depends on torch 1.6.0.

kornia: 0.4.0
torch: 1.6.0

Code

import torch
from byol_pytorch import BYOL
from torchvision import models

resnet = models.resnet50(pretrained=True)

learner = BYOL(
    resnet,
    image_size = 256,
    hidden_layer = 'avgpool'
)

Error

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-5-f272192ae7d5> in <module>
      8     resnet,
      9     image_size = 256,
---> 10     hidden_layer = 'avgpool'
     11 )

/opt/conda/lib/python3.7/site-packages/byol_pytorch/byol_pytorch.py in __init__(self, net, image_size, hidden_layer, projection_size, projection_hidden_size, augment_fn, moving_average_decay)
    157             RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1),
    158             augs.RandomResizedCrop((image_size, image_size)),
--> 159             color.Normalize(mean=torch.tensor([0.485, 0.456, 0.406]), std=torch.tensor([0.229, 0.224, 0.225]))
    160         )
    161 

AttributeError: module 'kornia.color' has no attribute 'Normalize'

Short-Term Solution

Is there a way to run the sample code without kornia dependency, e.g. by passing custom normalization and augmentations?

[lucidrains]

https://github.com/lucidrains/byol-pytorch/releases/tag/0.2.0 should be fixed!