Does _setup_transform() in Engine.fit use the train transforms properly?

[CarlosNacher]

Hi,

I was looking the Engine.fit() method and I realised that maybe we are incorrectly using eval_transform for train instead of train_transform; or maybe I have misunderstood the code structure. I mean:

If you look into the fit() method of the Engine (engine.py line 441 ) the method _setup_transform() is called. And that method (line 311 of engine.py) does:

# get transform
if datamodule and datamodule.transform:
    # a transform passed explicitly to the datamodule takes precedence
    transform = datamodule.transform

And if you go to datamodule (datamodule.py line 232 in the transform property):

def transform(self) -> Transform:
    """Property that returns the user-specified transform for the datamodule, if any.

    This property is accessed by the engine to set the transform for the model. The eval_transform takes precedence
    over the train_transform, because the transform that we store in the model is the one that should be used during
    inference.
    """
    if self._eval_transform:
        return self._eval_transform
    if self._train_transform:
        return self._train_transform
    return None

it seems that even in fit method, the eval_transform takes precedence over train_transform. But if you are fitting, you want train_transform, not eval one.

Am I misunderstanding anything?

Thank you so much in advance!