AttributeError: module 'albumentations' has no attribute 'Lambda'の解決法

google colabo上で下のpythonコードを実行したところ、

AttributeError: module 'albumentations' has no attribute 'Lambda'

と表示されます。古いversionで対応されていない、もしくは新しいversionで対応してないというのが原因だと思うのですが、古いやつをアンインストールして、対応するversionをインストールするコマンドが知りたいです。

同じ問題を抱えて解決された方がいらしたら教えていただきたいです。

試したこと
・albumentationの更新
!pip install -U albumentations>=0.3.0 --user
!pip install -U git+https://github.com/albu/albumentations --no-cache-dir

・albumentationのversionの確認
import albumentations
import torch
print(albumentations.version)

0.1.12

python
1train_dir = 'train'
2val_dir = 'val'
3
4preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
5
6# データセットを作成
7train_dataset = Dataset(
8    os.path.join(train_dir, 'images'),
9    os.path.join(train_dir, 'masks'),
10    augmentation=get_training_augmentation(),
11    preprocessing=get_preprocessing(preprocessing_fn),
12    classes=['cow1', 'cow2'],
13)
14valid_dataset = Dataset(
15    os.path.join(val_dir, 'images'),
16    os.path.join(val_dir, 'masks'),
17    preprocessing=get_preprocessing(preprocessing_fn),
18    classes=['cow1', 'cow2'],
19)
20
21train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True, num_workers=0)
22valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=False, num_workers=0)