Method train

torchvision/models/feature_extraction.py:320–332

Swap out the graph depending on the selected training mode. NOTE this should be safe when calling model.eval() because that just calls this with mode == False.

(self, mode=True)

            self._tracer_cls = self.graph._tracer_cls

    def train(self, mode=True):
        """
        Swap out the graph depending on the selected training mode.
        NOTE this should be safe when calling model.eval() because that just
        calls this with mode == False.
        """
        # NOTE: Only set self.graph if the current graph is not the desired
        # one. This saves us from recompiling the graph where not necessary.
        if mode and not self.training:
            self.graph = self.train_graph
        elif not mode and self.training:
            self.graph = self.eval_graph
        return super().train(mode=mode)

    def _deepcopy_init(self):
        # See __deepcopy__ below