Update predict.py

yahuiwei123 · Dec 21, 2023 · 1bf5e20 · 1bf5e20
1 parent 1e1d956
commit 1bf5e20
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/predict.py b/predict.py
@@ -40,7 +40,7 @@
 
     net = SegmentModel(output_size=output_size, out_cls=21, node=BiasLIFNode, step=step)
     # load model
-    torch.load(net, device, './checkpoints/Segment_SNN.pth')
+    net.load_state_dict(torch.load('./checkpoints/Segment_SNN.pth'))
     net = net.to(device)
     with torch.no_grad():
         for (idx, img) in enumerate(test_iter):