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):
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_2, Y_np2, pool_indices2 = self.encoderBlock2.forward(
Y_encoder_1, Y_np1, _ = self.encoderBlock1.forward(X)
Y_encoder_2, Y_np2, _ = self.encoderBlock2.forward(
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)
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)
del Y_encoder_3
......@@ -120,24 +120,24 @@ class BrainMapperUNet3D(nn.Module):
del Y_encoder_4
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_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_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_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)
......
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