Merge pull request #26 from unimusic/patch-3

Update unetSeg.py
UCSB-VRL · Mar 4, 2021 · 953694a · 953694a
2 parents 1800a54 + 6795d20
commit 953694a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/source/modules/nphprediction/source/unetSeg.py b/source/modules/nphprediction/source/unetSeg.py
@@ -205,9 +205,9 @@ def forward(self, x, prev):
 	net.seg3 = nn.Conv3d(128, num_classes, kernel_size=(1,1,1), stride=(1,1,1))
 	net.seg2 = nn.Conv3d(64, num_classes, kernel_size=(1,1,1), stride=(1,1,1))
 	net.seg1 = nn.Conv3d(32, num_classes, kernel_size=(1,1,1), stride=(1,1,1))
-	net.seg3.weight = nn.Parameter(unet.seg3.weight)
-	net.seg2.weight = nn.Parameter(unet.seg2.weight)
-	net.seg1.weight = nn.Parameter(unet.seg1.weight)
+	net.seg3.weight = nn.Parameter(net.seg3.weight)
+	net.seg2.weight = nn.Parameter(net.seg2.weight)
+	net.seg1.weight = nn.Parameter(net.seg1.weight)
 
 	net.cpu()
 	if gpu: