i’m guessing you need to detach self.cov_alpha since it depends on parameters of self.K_w
detach
self.cov_alpha
self.K_w
That did it, thanks!