fix: use specific exceptions and idiomatic None comparison

src/stage2/transport/integrators.py

@@ -61,7 +61,7 @@ def __forward_fn(self):
 
         try:
             sampler = sampler_dict[self.sampler_type]
-        except:
+        except KeyError:
             raise NotImplementedError("Smapler type not implemented.")
 
         return sampler

src/stage2/transport/transport.py

@@ -194,7 +194,7 @@ def training_losses(
         - x1: datapoint
         - model_kwargs: additional arguments for the model
         """
-        if model_kwargs == None:
+        if model_kwargs is None:
             model_kwargs = {}
 
         t, x0, x1 = self.sample(x1)

src/train.py

@@ -220,11 +220,11 @@ def guidance_value(key: str, default: float) -> float:
     if args.compile:
         try:
             rae.encode = torch.compile(rae.encode)
-        except:
+        except Exception:
             print('RAE ENCODE compile meets error, falling back to no compile')
         try:
             model.forward = torch.compile(model.forward)
-        except:
+        except Exception:
             print('MODEL FORWARD compile meets error, falling back to no compile')
     else:
         raise NotImplementedError('ARGS>COMPILE')