Pickle error and OOM when upgrading to 1.2.0

When upgrading from 1.1.6 to 1.2.0, I notices 2 changes

• Significant increase of gpu memory
• Pickle error in my module class (the object metric_fn is sure not pickable but the same code worked fine in 1.1.6)

Do you have an idea what changes in 1.2.0 may cause the issues ? Any suggestion for the memory problem ?

My pseudo code

class ClfModule(pl.LightningModule):
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.metric_fn = not_pickable_object()

I am using Huggingface transformers and datasets

[angadkalra]

@phihung My OOM went away with upgrade to v1.2.1 and using DDP. Please try.

[tchaton]

Dear @angadkalra,

Were using using AMP ? In 1.2.0, there were a bug and AMP wasn’t applied properly. It is solved in 1.2.1. Which could explain your OOM in 1.2.0 and not in 1.2.1.

Best, T.C