Commit 98e87284 authored by Andrei-Claudiu Roibu's avatar Andrei-Claudiu Roibu 🖥
Browse files

deleted pool_indices - not needed for upconv and leaking memory

parent 46d932ae
...@@ -99,18 +99,18 @@ class BrainMapperUNet3D(nn.Module): ...@@ -99,18 +99,18 @@ class BrainMapperUNet3D(nn.Module):
probability_map (torch.tensor): Output forward passed tensor through the U-net block probability_map (torch.tensor): Output forward passed tensor through the U-net block
""" """
Y_encoder_1, Y_np1, pool_indices1 = self.encoderBlock1.forward(X) Y_encoder_1, Y_np1, _ = self.encoderBlock1.forward(X)
Y_encoder_2, Y_np2, pool_indices2 = self.encoderBlock2.forward( Y_encoder_2, Y_np2, _ = self.encoderBlock2.forward(
Y_encoder_1) Y_encoder_1)
del Y_encoder_1 del Y_encoder_1
Y_encoder_3, Y_np3, pool_indices3 = self.encoderBlock3.forward( Y_encoder_3, Y_np3, _ = self.encoderBlock3.forward(
Y_encoder_2) Y_encoder_2)
del Y_encoder_2 del Y_encoder_2
Y_encoder_4, Y_np4, pool_indices4 = self.encoderBlock4.forward( Y_encoder_4, Y_np4, _ = self.encoderBlock4.forward(
Y_encoder_3) Y_encoder_3)
del Y_encoder_3 del Y_encoder_3
...@@ -120,24 +120,24 @@ class BrainMapperUNet3D(nn.Module): ...@@ -120,24 +120,24 @@ class BrainMapperUNet3D(nn.Module):
del Y_encoder_4 del Y_encoder_4
Y_decoder_1 = self.decoderBlock1.forward( Y_decoder_1 = self.decoderBlock1.forward(
Y_bottleNeck, Y_np4, pool_indices4) Y_bottleNeck, Y_np4)
del Y_bottleNeck, Y_np4, pool_indices4 del Y_bottleNeck, Y_np4
Y_decoder_2 = self.decoderBlock2.forward( Y_decoder_2 = self.decoderBlock2.forward(
Y_decoder_1, Y_np3, pool_indices3) Y_decoder_1, Y_np3)
del Y_decoder_1, Y_np3, pool_indices3 del Y_decoder_1, Y_np3
Y_decoder_3 = self.decoderBlock3.forward( Y_decoder_3 = self.decoderBlock3.forward(
Y_decoder_2, Y_np2, pool_indices2) Y_decoder_2, Y_np2)
del Y_decoder_2, Y_np2, pool_indices2 del Y_decoder_2, Y_np2
Y_decoder_4 = self.decoderBlock4.forward( Y_decoder_4 = self.decoderBlock4.forward(
Y_decoder_3, Y_np1, pool_indices1) Y_decoder_3, Y_np1)
del Y_decoder_3, Y_np1, pool_indices1 del Y_decoder_3, Y_np1
probability_map = self.classifier.forward(Y_decoder_4) probability_map = self.classifier.forward(Y_decoder_4)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment