Type casting of module input

I have a torch.nn.module implementation whose forward pass accepts as input a custom class structure that is inherited from a dict. All values of this object are torch Tensors.

The problem is that when I wrap my model with amp.initialize() and use amp level 02, the input which my model receives is a standard dict.

I managed to trace this issue to line 203 of the _initialize.py script: https://github.com/NVIDIA/apex/blob/master/apex/amp/_initialize.py#L203

The signature of the function is def new_fwd(*args, **kwargs): which means that my input class is broken up into a standard dict and the custom structure is lost.

Can anyone suggest a way around this? Preferably without having to modify the apex code itself.

Start with apex/amp/_initialize.py at line 203 and reproduce the report using a dict subclass whose values are torch Tensors, then wrap the model with amp.initialize() at amp level O2. Compare the input type received by forward before and after wrapping. Done means establishing whether the custom structure can be preserved or identifying a documented workaround that avoids modifying Apex.