fixing mapping of checkpoint

app.py

@@ -20,10 +20,10 @@ st.text(resp.status_code)
 pv.set_plot_theme("document")
 pl = pv.Plotter(shape=(1, 1),
                      window_size=(800, 800))
-print(torch.load("berea_generator_epoch_24.pth"))
+
 netG = DCGAN3D_G(64, 512, 1, 32, 1)
-netG.load_state_dict(torch.load("berea_generator_epoch_24.pth"))
-z = torch.randn(1, 512, 5, 5, 5)
+netG.load_state_dict(torch.load("berea_generator_epoch_24.pth", map_location=torch.device('cpu')))
+z = torch.randn(1, 512, 3, 3, 3)
 with torch.no_grad():
     X = netG(z)