class StrCategorize(Transform): def __init__(self, vocab): self.vocab = vocab self.s2i = {s:i for i,s in enumerate(vocab)} self.i2s = {i:s for i,s in enumerate(vocab)} def encodes(self, s:str): return self.s2i[s] def decodes(self, i:int): return self.i2s[i] transforms_ft = Pipeline([ load_img_and_label, Resize(256,method="squish"), Resize(224,method="crop"), ToTensor(), IntToFloatTensor(), Normalize.from_stats(*imagenet_stats, cuda=False), StrCategorize(vocab=['husky','wolf']), # <-- Transform is just for the target label ]) out = transforms_ft(fpath) print((out[0][0,:2,:2,:2], out[1]))