function signature patch and imports

rsk97 · Mar 9, 2023 · 6db0581 · 6db0581
1 parent 60acb2f
commit 6db0581
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/MoG.py b/MoG.py
@@ -43,8 +43,8 @@
 
         if i%5000 == 0:
             if method == 'ConsOpt':
-                plot_eigens(i)
-            plot_kde(i)
+                plot_eigens(i, gen_net, disc_net, params, gamma, path, device)
+            plot_kde(i,  method, sigma, gen_net, path, device, batch_size, z_dim, real_input=False)
 
             gen_path = join(path, 'Models', 'gen_' + method + "_" + str(i) + '.pt')
             disc_path = join(path, 'Models', 'disc_'+ method + "_" + str(i) + '.pt')

diff --git a/utils/plot.py b/utils/plot.py
@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 from os.path import join
 from torch import autograd
 from utils.utils import batch_net_outputs, net_losses, complex_scatter_plot, kde
@@ -38,7 +39,7 @@ def plot_eigens(iteration, gen, disc, params, gamma, path, device):
     complex_scatter_plot(eigens2.cpu().detach().numpy(), bbox=[-1.0, 1.0, -0.15, 0.15], save_file=save_path, cmap=cmap)
 
 
-def plot_kde(iteration, method, sigma, gen, disc, params, gamma, path, device, batch_size, z_dim, real_input=False):
+def plot_kde(iteration, method, sigma, gen, path, device, batch_size, z_dim, real_input=False):
     z = torch.normal(mean=0, std=1, size=[batch_size * 5, z_dim]).to(device)
     inp = gen(z)