flamo 0.1.13__tar.gz → 0.2.0__tar.gz

{flamo-0.1.13 → flamo-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flamo
-Version: 0.1.13
+Version: 0.2.0
 Summary: An Open-Source Library for Frequency-Domain Differentiable Audio Processing
 Project-URL: Homepage, https://github.com/gdalsanto/flamo
 Project-URL: Issues, https://github.com/gdalsanto/flamo/issues

{flamo-0.1.13 → flamo-0.2.0}/examples/e0_siso.py

@@ -20,8 +20,8 @@ def example_fft(args) -> None:
     Use of FFT and iFFT modules.
     """
     # ------------------ Module Definition ------------------
-    fft = dsp.FFT(nfft=args.nfft)
-    ifft = dsp.iFFT(nfft=args.nfft)
+    fft = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    ifft = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     # ------------------ Signal Definition ------------------
     x = signal_gallery(
@@ -31,6 +31,7 @@ def example_fft(args) -> None:
         n=1,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # ------------------ Apply FFT and iFFT -----------------
@@ -58,9 +59,9 @@ def example_gains(args) -> None:
 
     # ------------------- DSP Definition --------------------
     channels = 1
-    filter = dsp.parallelGain(size=(channels,), nfft=args.nfft, device=args.device)
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    filter = dsp.parallelGain(size=(channels,), nfft=args.nfft, device=args.device, dtype=args.dtype)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     my_dsp = nn.Sequential(input_layer, filter, output_layer)
 
@@ -74,6 +75,7 @@ def example_gains(args) -> None:
         n=channels,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Apply filter
@@ -104,9 +106,9 @@ def example_gains_2(args) -> None:
     # ------------------- DSP Definition -------------------
     in_ch = 1
     out_ch = 1
-    filter = dsp.Gain(size=(out_ch, in_ch), nfft=args.nfft, device=args.device)
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    filter = dsp.Gain(size=(out_ch, in_ch), nfft=args.nfft, device=args.device, dtype=args.dtype)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     my_dsp = nn.Sequential(input_layer, filter, output_layer)
 
@@ -120,6 +122,7 @@ def example_gains_2(args) -> None:
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Apply filter before changes
@@ -168,9 +171,10 @@ def example_fir(args) -> None:
         nfft=args.nfft,
         requires_grad=True,
         device=args.device,
+        dtype=args.dtype,
     )
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     model = nn.Sequential(input_layer, filter, output_layer)
 
@@ -183,6 +187,7 @@ def example_fir(args) -> None:
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Target impulse response
@@ -194,6 +199,7 @@ def example_fir(args) -> None:
         rate=2,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Dataset
@@ -252,6 +258,7 @@ if __name__ == "__main__":
     # ---------------------- Processing -------------------
     parser.add_argument("--nfft", type=int, default=96000, help="FFT size")
     parser.add_argument("--samplerate", type=int, default=48000, help="sampling rate")
+    parser.add_argument("--dtype", type=str, default="float64", choices=["float32", "float64"], help="data type for tensors")
     # ----------------------- Dataset ----------------------
     parser.add_argument(
         "--batch_size", type=int, default=1, help="batch size for training"
@@ -288,6 +295,9 @@ if __name__ == "__main__":
     if args.device == "cuda" and not torch.cuda.is_available():
         args.device = "cpu"
 
+    # convert dtype string to torch dtype
+    args.dtype = torch.float32 if args.dtype == "float32" else torch.float64
+
     # make output directory
     if args.train_dir is not None:
         if not os.path.isdir(args.train_dir):

{flamo-0.1.13 → flamo-0.2.0}/examples/e1_mimo.py

@@ -36,6 +36,7 @@ def example_matrix(args) -> None:
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
     
     # Test different matrix types
@@ -54,10 +55,11 @@ def example_matrix(args) -> None:
             matrix_type=matrix_type,
             nfft=args.nfft,
             device=args.device,
+            dtype=args.dtype,
         )
         
-        input_layer = dsp.FFT(nfft=args.nfft)
-        output_layer = dsp.iFFT(nfft=args.nfft)
+        input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+        output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
         
         my_dsp = nn.Sequential(input_layer, matrix_filter, output_layer)
         
@@ -94,9 +96,10 @@ def example_delays(args) -> None:
         nfft=args.nfft,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     my_dsp = nn.Sequential(input_layer, filter, output_layer)
 
@@ -110,6 +113,7 @@ def example_delays(args) -> None:
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Apply filter
@@ -147,9 +151,10 @@ def example_biquads(args) -> None:
         fs=args.samplerate,
         requires_grad=True,
         device=args.device,
+        dtype=args.dtype,
     )
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.Transform(lambda x: torch.abs(x))
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.Transform(lambda x: torch.abs(x), dtype=args.dtype)
 
     model = nn.Sequential(input_layer, filter, output_layer)
 
@@ -163,17 +168,18 @@ def example_biquads(args) -> None:
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Target frequency responses
     f_cut_1 = 500  # Cut-off frequency for the first lowpass filter
     g_1 = 10  # Bandpass gain for the first lowpass filter
-    b_lp_1, a_lp_1 = lowpass_filter(f_cut_1, g_1, args.samplerate, device=args.device)
+    b_lp_1, a_lp_1 = lowpass_filter(f_cut_1, g_1, args.samplerate, device=args.device, dtype=args.dtype)
     H_lp_1 = biquad2tf(b=b_lp_1, a=a_lp_1, nfft=args.nfft)
 
     f_cut_2 = 5000  # Cut-off frequency for the second lowpass filter
     g_2 = 0.7  # Bandpass gain for the second lowpass filter
-    b_lp_2, a_lp_2 = lowpass_filter(f_cut_2, g_2, args.samplerate, device=args.device)
+    b_lp_2, a_lp_2 = lowpass_filter(f_cut_2, g_2, args.samplerate, device=args.device, dtype=args.dtype)
     H_lp_2 = biquad2tf(b=b_lp_2, a=a_lp_2, nfft=args.nfft)
 
     target = torch.stack([torch.abs(H_lp_1), torch.abs(H_lp_2)], dim=1).unsqueeze(0)
@@ -242,6 +248,7 @@ if __name__ == "__main__":
     # ---------------------- Processing -------------------
     parser.add_argument("--nfft", type=int, default=96000, help="FFT size")
     parser.add_argument("--samplerate", type=int, default=48000, help="sampling rate")
+    parser.add_argument("--dtype", type=str, default="float64", choices=["float32", "float64"], help="data type for tensors")
     # ----------------------- Dataset ----------------------
     parser.add_argument(
         "--batch_size", type=int, default=1, help="batch size for training"
@@ -278,6 +285,9 @@ if __name__ == "__main__":
     if args.device == "cuda" and not torch.cuda.is_available():
         args.device = "cpu"
 
+    # convert dtype string to torch dtype
+    args.dtype = torch.float32 if args.dtype == "float32" else torch.float64
+
     # make output directory
     if args.train_dir is not None:
         if not os.path.isdir(args.train_dir):

{flamo-0.1.13 → flamo-0.2.0}/examples/e1_vn.py

@@ -5,6 +5,7 @@ This example shows how to use the new Filter-based velvet noise classes
 that are differentiable and can be used within the FLAMO processing chain.
 """
 
+import torch
 import torch.nn as nn   
 
 import matplotlib.pyplot as plt
@@ -21,6 +22,7 @@ def example_velvet_noise_filter():
     in_ch, out_ch = 1, 1
     nfft = 2048  # FFT size for the filter
     length = nfft # Length of the filter
+    dtype = torch.float64  # Can be changed to torch.float32 if needed
 
     # Create a velvet noise filter
     # size = (length, output_channels, input_channels)
@@ -28,11 +30,12 @@ def example_velvet_noise_filter():
         size=(length, out_ch, in_ch),  # 1024 samples, 1x1 matrix
         density=1000.0,     # 1000 impulses per second
         sample_rate=48000,
-        requires_grad=True  # Make it differentiable
+        requires_grad=True,  # Make it differentiable
+        dtype=dtype,
     )
     
-    input_layer = dsp.FFT(nfft=nfft)
-    output_layer = dsp.iFFT(nfft=nfft)
+    input_layer = dsp.FFT(nfft=nfft, dtype=dtype)
+    output_layer = dsp.iFFT(nfft=nfft, dtype=dtype)
 
     my_dsp = nn.Sequential(input_layer, velvet_filter, output_layer)
 
@@ -46,6 +49,7 @@ def example_velvet_noise_filter():
         n=in_ch,
         fs=48000,
         device="cpu",
+        dtype=dtype,
     )
 
     output_vn = my_dsp(input_sig)

{flamo-0.1.13 → flamo-0.2.0}/examples/e2_chaining_filters.py

@@ -33,8 +33,8 @@ def example_mimo(args):
         fs=args.samplerate,
         device=args.device,
     )
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     my_dsp = nn.Sequential(input_layer, filter1, filter2, output_layer)
 
@@ -48,6 +48,7 @@ def example_mimo(args):
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Apply filter
@@ -87,9 +88,10 @@ def example_siso(args):
         nfft=args.nfft,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype
     )
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     my_dsp = nn.Sequential(input_layer, filter1, filter2, output_layer)
 
@@ -103,6 +105,7 @@ def example_siso(args):
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Apply filter
@@ -143,9 +146,10 @@ def example_assign_new_values(args):
         nfft=args.nfft,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype
     )
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     my_dsp = nn.Sequential(input_layer, filter1, filter2, output_layer)
 
@@ -171,6 +175,7 @@ def example_assign_new_values(args):
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Apply filter
@@ -209,9 +214,10 @@ def example_requires_grad(args):
         nfft=args.nfft,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     model = nn.Sequential(input_layer, filter1, filter2, output_layer)
 
@@ -225,6 +231,7 @@ def example_requires_grad(args):
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Target
@@ -298,6 +305,7 @@ if __name__ == "__main__":
     # ---------------------- Processing -------------------
     parser.add_argument("--nfft", type=int, default=96000, help="FFT size")
     parser.add_argument("--samplerate", type=int, default=48000, help="sampling rate")
+    parser.add_argument("--dtype", type=str, default="float64", choices=["float32", "float64"], help="data type for tensors")
     # ----------------------- Dataset ----------------------
     parser.add_argument(
         "--batch_size", type=int, default=1, help="batch size for training"
@@ -334,6 +342,9 @@ if __name__ == "__main__":
     if args.device == "cuda" and not torch.cuda.is_available():
         args.device = "cpu"
 
+    # convert dtype string to torch dtype
+    args.dtype = torch.float32 if args.dtype == "float32" else torch.float64
+
     # make output directory
     if args.train_dir is not None:
         if not os.path.isdir(args.train_dir):

{flamo-0.1.13 → flamo-0.2.0}/examples/e3_series_class.py

@@ -50,8 +50,8 @@ def example_series(args):
         size=(out_ch, btw_ch), nfft=args.nfft, requires_grad=False, device=args.device
     )
     # Input and output layers
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     # Series class
     my_dsp = system.Series(input_layer, filter1, filter2, filter3, output_layer)
@@ -93,8 +93,8 @@ def example_series_with_error(args):
         device=args.device,
     )
     # Input and output layers
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     # Series class
     my_dsp = system.Series(input_layer, filter1, filter2, filter3, output_layer)
@@ -131,8 +131,8 @@ def example_series_OrderedDict(args):
         size=(out_ch, btw_ch), nfft=args.nfft, requires_grad=False, device=args.device
     )
     # Input and output layers
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     # Series class
     my_dsp = system.Series(
@@ -182,8 +182,8 @@ def example_series_nesting(args):
         size=(out_ch, btw_ch), nfft=args.nfft, requires_grad=False, device=args.device
     )
     # Input and output layers
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     # Series class
     filters = OrderedDict({"Gains": filter1, "Delays": filter2, "Eqs": filter3})
@@ -221,9 +221,10 @@ def example_series_training(args):
         nfft=args.nfft,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     # Series of filters
     filters = OrderedDict(
@@ -252,6 +253,7 @@ def example_series_training(args):
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Target
@@ -336,6 +338,7 @@ def example_series_utils(args):
         fs=args.samplerate,
         device=args.device,
         alias_decay_db=alias_decay_db,
+        dtype=args.dtype
     )
     feedback = dsp.Matrix(
         size=(channel_n, channel_n),
@@ -343,6 +346,7 @@ def example_series_utils(args):
         matrix_type="orthogonal",
         device=args.device,
         alias_decay_db=alias_decay_db,
+        dtype=args.dtype
     )
     feedback_loop = system.Recursion(fF=feedforward, fB=feedback)
 
@@ -350,13 +354,14 @@ def example_series_utils(args):
     my_dsp = system.Series(OrderedDict({"Recursion": feedback_loop}))
 
     # New filters to add an the beginning
-    input_layer = dsp.FFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
     input_gains = dsp.parallelGain(
         size=(channel_n,),
         nfft=args.nfft,
         requires_grad=False,
         device=args.device,
         alias_decay_db=alias_decay_db,
+        dtype=args.dtype
     )
     # New filter to add in the middle
     output_gains = dsp.Gain(
@@ -365,6 +370,7 @@ def example_series_utils(args):
         requires_grad=False,
         device=args.device,
         alias_decay_db=alias_decay_db,
+        dtype=args.dtype
     )
     # New filters to add at the end
     equalization = dsp.GEQ(
@@ -373,8 +379,9 @@ def example_series_utils(args):
         requires_grad=False,
         device=args.device,
         alias_decay_db=alias_decay_db,
+        dtype=args.dtype
     )
-    output_layer = dsp.iFFTAntiAlias(nfft=args.nfft, alias_decay_db=alias_decay_db)
+    output_layer = dsp.iFFTAntiAlias(nfft=args.nfft, alias_decay_db=alias_decay_db, dtype=args.dtype)
 
     # DSP so far
     print(my_dsp)
@@ -409,6 +416,7 @@ def example_series_utils(args):
         n=channel_n,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
     y = my_dsp(input_signal)
 
@@ -436,6 +444,7 @@ if __name__ == "__main__":
     # ---------------------- Processing -------------------
     parser.add_argument("--nfft", type=int, default=96000, help="FFT size")
     parser.add_argument("--samplerate", type=int, default=48000, help="sampling rate")
+    parser.add_argument("--dtype", type=str, default="float64", choices=["float32", "float64"], help="data type for tensors")
     # ----------------------- Dataset ----------------------
     parser.add_argument(
         "--batch_size", type=int, default=1, help="batch size for training"
@@ -472,6 +481,9 @@ if __name__ == "__main__":
     if args.device == "cuda" and not torch.cuda.is_available():
         args.device = "cpu"
 
+    # convert dtype string to torch dtype
+    args.dtype = torch.float32 if args.dtype == "float32" else torch.float64
+
     # make output directory
     if args.train_dir is not None:
         if not os.path.isdir(args.train_dir):

{flamo-0.1.13 → flamo-0.2.0}/examples/e4_recursion.py

@@ -30,8 +30,9 @@ def example_recursion(args):
         nfft=args.nfft,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
-    attenuation = dsp.parallelGain(size=(out_ch,), nfft=args.nfft, device=args.device)
+    attenuation = dsp.parallelGain(size=(out_ch,), nfft=args.nfft, device=args.device, dtype=args.dtype)
     rand_vector = torch.rand(attenuation.param.shape)
     attenuation.assign_value(0.3 * rand_vector / torch.norm(rand_vector, p=2))
     feedforward_path = OrderedDict({"delays": delays, "attenuation": attenuation})
@@ -42,6 +43,7 @@ def example_recursion(args):
         matrix_type="orthogonal",
         nfft=args.nfft,
         device=args.device,
+        dtype=args.dtype,
     )
 
     feedback_path = OrderedDict({"feedback_matrix": feedback_matrix})
@@ -50,8 +52,8 @@ def example_recursion(args):
     recursion = system.Recursion(fF=feedforward_path, fB=feedback_path)
 
     # Input and output layers
-    input_layer = dsp.FFT(nfft=args.nfft)
-    output_layer = dsp.iFFT(nfft=args.nfft)
+    input_layer = dsp.FFT(nfft=args.nfft, dtype=args.dtype)
+    output_layer = dsp.iFFT(nfft=args.nfft, dtype=args.dtype)
 
     my_dsp = system.Series(
         OrderedDict(
@@ -73,6 +75,7 @@ def example_recursion(args):
         n=in_ch,
         fs=args.samplerate,
         device=args.device,
+        dtype=args.dtype,
     )
 
     # Apply filter
@@ -104,6 +107,7 @@ if __name__ == "__main__":
     # ---------------------- Processing -------------------
     parser.add_argument("--nfft", type=int, default=96000, help="FFT size")
     parser.add_argument("--samplerate", type=int, default=48000, help="sampling rate")
+    parser.add_argument("--dtype", type=str, default="float64", choices=["float32", "float64"], help="data type for tensors")
     # ----------------------- Dataset ----------------------
     parser.add_argument(
         "--batch_size", type=int, default=1, help="batch size for training"
@@ -141,6 +145,9 @@ if __name__ == "__main__":
     if args.device == "cuda" and not torch.cuda.is_available():
         args.device = "cpu"
 
+    # convert dtype string to torch dtype
+    args.dtype = torch.float32 if args.dtype == "float32" else torch.float64
+
     # make output directory
     if args.train_dir is not None:
         if not os.path.isdir(args.train_dir):

{flamo-0.1.13 → flamo-0.2.0}/examples/e4_recursion_nn.py

@@ -214,8 +214,8 @@ class nnComb(nn.Module):
         comb = system.Recursion(fF=delays, fB=filt)
 
         # Create the model with Shell
-        input_layer = dsp.FFT(args.nfft)
-        output_layer = dsp.Transform(transform=lambda x: torch.abs(x))
+        input_layer = dsp.FFT(args.nfft, dtype=args.dtype)
+        output_layer = dsp.Transform(transform=lambda x: torch.abs(x), dtype=args.dtype)
         self.comb = system.Shell(
             core=comb, input_layer=input_layer, output_layer=output_layer
         )
@@ -344,6 +344,7 @@ if __name__ == "__main__":
 
     parser.add_argument("--nfft", type=int, default=32000, help="FFT size")
     parser.add_argument("--samplerate", type=int, default=16000, help="sampling rate")
+    parser.add_argument("--dtype", type=str, default="float64", choices=["float32", "float64"], help="data type for tensors")
     parser.add_argument("--num", type=int, default=2**10, help="dataset size")
     parser.add_argument(
         "--device", type=str, default="cuda", help="device to use for computation"
@@ -373,6 +374,9 @@ if __name__ == "__main__":
     if args.device == "cuda" and not torch.cuda.is_available():
         args.device = "cpu"
 
+    # convert dtype string to torch dtype
+    args.dtype = torch.float32 if args.dtype == "float32" else torch.float64
+
     # Make output directory
     if args.train_dir is not None:
         if not os.path.isdir(args.train_dir):