Hello i need help with measuring Iou. Can someone help me or give example how can i save best model during the training of this model ?
class Model(torch.nn.Module):

    def __init__(self, input_shape= (3, 448, 640), output_shape= (2, 448, 640)):
        super(Model, self).__init__()

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.layers_encoder_0 = [
            self.conv_bn(input_shape[0], 64, 2),

            self.conv_bn(64, 64, 1),
            self.conv_bn(64, 128, 2),

            self.conv_bn(128, 128, 1),
            self.conv_bn(128, 256, 2)
        ]

        self.layers_encoder_1 = [
            self.conv_bn(256, 256, 1),
            self.conv_bn(256, 256, 2),

            self.conv_bn(256, 256, 1),
            self.conv_bn(256, 256, 2)
        ]

        self.layers_decoder = [
            self.conv_bn(256 + 256, 256, 1),
            self.conv_bn(256, 128, 1),
            self.conv_bn(128, 128, 1),

            nn.Conv2d(128, output_shape[0], kernel_size=1, stride=1, padding=0),
            nn.Upsample(scale_factor=8, mode="bilinear", align_corners=False)
        ]

        for i in range(len(self.layers_encoder_0)):
            if hasattr(self.layers_encoder_0[i], "weight"):
                torch.nn.init.xavier_uniform_(self.layers_encoder_0[i].weight)

        for i in range(len(self.layers_encoder_1)):
            if hasattr(self.layers_encoder_1[i], "weight"):
                torch.nn.init.xavier_uniform_(self.layers_encoder_1[i].weight)

        for i in range(len(self.layers_decoder)):
            if hasattr(self.layers_decoder[i], "weight"):
                torch.nn.init.xavier_uniform_(self.layers_decoder[i].weight)

        self.model_encoder_0 = nn.Sequential(*self.layers_encoder_0)
        self.model_encoder_0.to(self.device)

        self.model_encoder_1 = nn.Sequential(*self.layers_encoder_1)
        self.model_encoder_1.to(self.device)

        self.model_decoder = nn.Sequential(*self.layers_decoder)
        self.model_decoder.to(self.device)

        print(self.model_encoder_0)
        print(self.model_encoder_1)
        print(self.model_decoder)

    def forward(self, x):
        encoder_0 = self.model_encoder_0(x)
        encoder_1 = self.model_encoder_1(encoder_0)

        encoder_1_up = torch.nn.functional.interpolate(encoder_1, scale_factor=4, mode="nearest")

        d_in = torch.cat([encoder_0, encoder_1_up], dim=1)

        y = self.model_decoder(d_in)
        return y

    def save(self, path):
        torch.save(self.model_encoder_0.state_dict(), path + "model_encoder_0.pt")
        torch.save(self.model_encoder_1.state_dict(), path + "model_encoder_1.pt")
        torch.save(self.model_decoder.state_dict(), path + "model_decoder.pt")

    def load(self, path):
        self.model_encoder_0.load_state_dict(torch.load(path + "model_encoder_0.pt", map_location=self.device))
        self.model_encoder_1.load_state_dict(torch.load(path + "model_encoder_1.pt", map_location=self.device))
        self.model_decoder.load_state_dict(torch.load(path + "model_decoder.pt", map_location=self.device))

        self.model_encoder_0.eval()
        self.model_encoder_1.eval()
        self.model_decoder.eval()

    def conv_bn(self, inputs, outputs, stride):
        return nn.Sequential(
            nn.Conv2d(inputs, outputs, kernel_size=3, stride=stride, padding=1),
            nn.BatchNorm2d(outputs),
            nn.ReLU(inplace=True))

    def tconv_bn(self, inputs, outputs, stride):
        return nn.Sequential(
            nn.ConvTranspose2d(inputs, outputs, kernel_size=3, stride=stride, padding=1, output_padding=1),
            nn.BatchNorm2d(outputs),
            nn.ReLU(inplace=True))