Custom head functionality removed from ImageDataBunch.from_folder?

Edit: I realized that I can just add the custom head when I declare the learner. So I no longer need assistance.

Hi, I am noticing that some code I had been running successfully using previous versions of FastAI has apparently been broken by recent updates.

I was loading data as a bunch into a folder and adding a custom head. Please tell me how to do the same thing in the current version of FastAI.

The code is:

head = nn.Sequential(
Flatten(),
nn.Dropout(0.5),
nn.Linear(32768,256),
nn.ReLU(),
nn.BatchNorm1d(256),
nn.Dropout(0.5),
nn.Linear(256,7)
)

data = ImageDataBunch.from_folder(path, ds_tfms=get_transforms(), size=sz, custom_head=head, bs=64).normalize(imagenet_stats)

This used to run fine, but when I try now, it gives the following output:

/home/paperspace/fastai/fastai/basic_data.py:207: UserWarning: There seems to be something wrong with your dataset, can’t access self.train_ds[i] for all i in [351, 299, 378, 13, 381, 161, 12, 158, 99, 16, 14, 348, 148, 117, 79, 290, 207, 278, 118, 202, 127, 265, 43, 230, 112, 312, 157, 277, 262, 7, 15, 1, 114, 263, 27, 325, 258, 186, 177, 184, 120, 135, 91, 2, 51, 23, 180, 198, 193, 72, 38, 57, 357, 349, 89, 126, 150, 121, 46, 156, 185, 140, 143, 62]
warn(f"There seems to be something wrong with your dataset, can’t access self.train_ds[i] for all i in {idx}")

And when I try to do data.show_batch(), it gives the following:

TypeError: Traceback (most recent call last):
File “/home/paperspace/anaconda3/envs/fastai/lib/python3.6/site-packages/torch/utils/data/dataloader.py”, line 138, in _worker_loop
samples = collate_fn([dataset[i] for i in batch_indices])
File “/home/paperspace/anaconda3/envs/fastai/lib/python3.6/site-packages/torch/utils/data/dataloader.py”, line 138, in
samples = collate_fn([dataset[i] for i in batch_indices])
File “/home/paperspace/fastai/fastai/data_block.py”, line 566, in getitem
x = x.apply_tfms(self.tfms, **self.tfmargs)
TypeError: apply_tfms() got an unexpected keyword argument ‘head’

I am also facing the same issue with latest version of fastai

Version: 1.0.40.dev0