Wrong log prob example behavior is fixed

README.md

@@ -47,6 +47,5 @@ python main.py --model=cnf --sample_dataset=circles
 
 # Sample Results
 
-![cnf-viz](https://user-images.githubusercontent.com/72425253/126116823-a014f13a-1171-4309-898f-0b6aedd84649.gif)
-![cnf-viz](https://user-images.githubusercontent.com/72425253/126117205-fa68c16b-fba1-48a0-a965-3ac6cb5e201c.gif)
-
+![cnf-viz](https://user-images.githubusercontent.com/72425253/126124351-44e00438-055e-4b1c-90ee-758a545dd602.gif)
+![cnf-viz](https://user-images.githubusercontent.com/72425253/126124648-dcb3f8f4-396a-447c-96cf-f9304377fa48.gif)

cnf_torch.py

@@ -15,7 +15,7 @@ import torch.optim as optim
 parser = argparse.ArgumentParser()
 parser.add_argument('--adjoint', action='store_true')
 parser.add_argument('--viz', action='store_true', default=True)
-parser.add_argument('--niters', type=int, default=10000)
+parser.add_argument('--niters', type=int, default=1)
 parser.add_argument('--lr', type=float, default=1e-3)
 parser.add_argument('--num_samples', type=int, default=512)
 parser.add_argument('--width', type=int, default=64)
@@ -148,9 +148,9 @@ if __name__ == '__main__':
 
     # model
     func = CNF(in_out_dim=2, hidden_dim=args.hidden_dim, width=args.width).to(device)
-    for param in func.parameters():
-        nn.init.constant_(param, 0.1)
-    func(torch.tensor(0.).to(device), (torch.tensor([[0., 1.], [2., 3.], [4., 5.]]).to(device), torch.zeros((2, 1)).to(device)))
+    # for param in func.parameters():
+    #     nn.init.constant_(param, 0.1)
+    # func(torch.tensor(0.).to(device), (torch.tensor([[0., 1.], [2., 3.], [4., 5.]]).to(device), torch.zeros((2, 1)).to(device)))
     optimizer = optim.Adam(func.parameters(), lr=args.lr)
     p_z0 = torch.distributions.MultivariateNormal(
         loc=torch.tensor([0.0, 0.0]).to(device),

main.py

@@ -22,7 +22,7 @@ if __name__ == '__main__':
     parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
     parser.add_argument("--n_epoch", type=int, default=10, help="Total number of epoch")
     parser.add_argument("--batch_size", type=int, default=32, help="Number of images in batch")
-    parser.add_argument("--sample_dataset", type=str, choices=['circels', 'moons'], default="circles",
+    parser.add_argument("--sample_dataset", type=str, choices=['circles', 'moons'], default="circles",
                         help="Sample dataset")
     parser.add_argument("--viz", action='store_true')