How to determine if the model is in train or eval phase?

1

I’m writing a network where one parameter (I used register_buffer) is expected to be updated if the model is under model.train(), and not to change if is under model.eval(). So I need a (boolean) criterion such that in the forward function, the parameter will be updated based on whether the model is under train or eval. Is there any function to get this value? Thanks!

2

Use the training attribute. its True when in train mode and False when in eval mode.

2 Likes
4

Thanks. I used “if self.training is True”, and it worked!