Fixing: Torch 2.0 fails with cuda=False (sdv-dev#291)

amontanez24 · web-flow · commit a20fad092f88 · 2023-05-11T15:31:57.000-07:00
* Making changes to get torch 2.0 to pass demo

* updating branch

* handling nans from gumbel_softmax
diff --git a/ctgan/synthesizers/ctgan.py b/ctgan/synthesizers/ctgan.py
@@ -5,7 +5,6 @@
 import numpy as np
 import pandas as pd
 import torch
-from packaging import version
 from torch import optim
 from torch.nn import BatchNorm1d, Dropout, LeakyReLU, Linear, Module, ReLU, Sequential, functional
 
@@ -197,15 +196,12 @@ def _gumbel_softmax(logits, tau=1, hard=False, eps=1e-10, dim=-1):
         Returns:
             Sampled tensor of same shape as logits from the Gumbel-Softmax distribution.
         """
-        if version.parse(torch.__version__) < version.parse('1.2.0'):
-            for i in range(10):
-                transformed = functional.gumbel_softmax(logits, tau=tau, hard=hard,
-                                                        eps=eps, dim=dim)
-                if not torch.isnan(transformed).any():
-                    return transformed
-            raise ValueError('gumbel_softmax returning NaN.')
+        for _ in range(10):
+            transformed = functional.gumbel_softmax(logits, tau=tau, hard=hard, eps=eps, dim=dim)
+            if not torch.isnan(transformed).any():
+                return transformed
 
-        return functional.gumbel_softmax(logits, tau=tau, hard=hard, eps=eps, dim=dim)
+        raise ValueError('gumbel_softmax returning NaN.')
 
     def _apply_activate(self, data):
         """Apply proper activation function to the output of the generator."""
@@ -381,7 +377,7 @@ def fit(self, train_data, discrete_columns=(), epochs=None):
                         real_cat, fake_cat, self._device, self.pac)
                     loss_d = -(torch.mean(y_real) - torch.mean(y_fake))
 
-                    optimizerD.zero_grad()
+                    optimizerD.zero_grad(set_to_none=False)
                     pen.backward(retain_graph=True)
                     loss_d.backward()
                     optimizerD.step()
@@ -412,7 +408,7 @@ def fit(self, train_data, discrete_columns=(), epochs=None):
 
                 loss_g = -torch.mean(y_fake) + cross_entropy
 
-                optimizerG.zero_grad()
+                optimizerG.zero_grad(set_to_none=False)
                 loss_g.backward()
                 optimizerG.step()
 
diff --git a/setup.py b/setup.py
@@ -18,8 +18,8 @@
     "pandas>=1.1.3;python_version<'3.10'",
     "pandas>=1.3.4;python_version>='3.10'",
     "scikit-learn>=1.1.3,<2;python_version>='3.10'",
-    "torch>=1.8.0,<2;python_version<'3.10'",
-    "torch>=1.11.0,<2;python_version>='3.10'",
+    "torch>=1.8.0;python_version<'3.10'",
+    "torch>=1.11.0;python_version>='3.10'",
     'rdt>=1.3.0,<2.0',
 ]
 
