diff --git a/ddsp/losses.py b/ddsp/losses.py
index be3828c..76dfbfd 100644
--- a/ddsp/losses.py
+++ b/ddsp/losses.py
@@ -135,8 +135,9 @@ class SpectralLoss(tfkl.Layer):
                                                      self.loss_type)
 
     if self.loudness_weight > 0:
-      target = spectral_ops.compute_loudness(target_audio, n_fft=2048)
-      value = spectral_ops.compute_loudness(audio, n_fft=2048)
+      target = spectral_ops.compute_loudness(target_audio, n_fft=2048,
+                                             use_tf=True)
+      value = spectral_ops.compute_loudness(audio, n_fft=2048, use_tf=True)
       loss += self.loudness_weight * mean_difference(target, value,
                                                      self.loss_type)
 
diff --git a/ddsp/spectral_ops.py b/ddsp/spectral_ops.py
index df9df08..b9466d0 100644
--- a/ddsp/spectral_ops.py
+++ b/ddsp/spectral_ops.py
@@ -155,10 +155,9 @@ def diff(x, axis=-1):
   begin_front = [0 for _ in range(len(shape))]
   begin_front[axis] = 1
 
-  size = shape.as_list()
-  size[axis] -= 1
-  slice_front = tf.slice(x, begin_front, size)
-  slice_back = tf.slice(x, begin_back, size)
+  shape[axis] -= 1
+  slice_front = tf.slice(x, begin_front, shape)
+  slice_back = tf.slice(x, begin_back, shape)
   d = slice_front - slice_back
   return d
 
diff --git a/ddsp/version.py b/ddsp/version.py
index 3c2b321..e714ae9 100644
--- a/ddsp/version.py
+++ b/ddsp/version.py
@@ -19,4 +19,4 @@ Stored in a separate file so that setup.py can reference the version without
 pulling in all the dependencies in __init__.py.
 """
 
-__version__ = '0.2.1'
+__version__ = '0.2.2'

All alternative spectral losses throw an error
SpectralLoss:
I am trying to use alternative loss weights, but all of them throw an error except logmag_weight and mag_weight:

```
INFO:tensorflow:Error reported to Coordinator: in user code:

    /usr/local/lib/python3.6/dist-packages/ddsp/training/models.py:64 __call__  *
        results = super().__call__(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/ddsp/training/models.py:123 call  *
        loss = loss_obj(features['audio'], audio_gen)
    /usr/local/lib/python3.6/dist-packages/ddsp/losses.py:107 call  *
        target = diff(target_mag, axis=1)
    /usr/local/lib/python3.6/dist-packages/ddsp/spectral_ops.py:158 diff  *
        size = shape.as_list()

    AttributeError: 'list' object has no attribute 'as_list'
```

or in case of trying to use loudness:
```
INFO:tensorflow:Error reported to Coordinator: in user code:

    /usr/local/lib/python3.6/dist-packages/ddsp/training/models.py:64 __call__  *
        results = super().__call__(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/ddsp/training/models.py:123 call  *
        loss = loss_obj(features['audio'], audio_gen)
    /usr/local/lib/python3.6/dist-packages/ddsp/losses.py:138 call  *
        target = spectral_ops.compute_loudness(target_audio, n_fft=2048)
    /usr/local/lib/python3.6/dist-packages/ddsp/spectral_ops.py:209 compute_loudness  *
        s = stft_fn(audio, frame_size=n_fft, overlap=overlap, pad_end=True)
    /usr/local/lib/python3.6/dist-packages/ddsp/spectral_ops.py:61 stft_np  *
        audio = np.pad(audio, padding, 'constant')
    <__array_function__ internals>:6 pad  **
        
    /usr/local/lib/python3.6/dist-packages/numpy/lib/arraypad.py:741 pad
        array = np.asarray(array)
    /usr/local/lib/python3.6/dist-packages/numpy/core/_asarray.py:85 asarray
        return array(a, dtype, copy=False, order=order)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py:749 __array__
        " array.".format(self.name))

    NotImplementedError: Cannot convert a symbolic Tensor (fn:0) to a numpy array.
```
diff --git a/ddsp/losses_test.py b/ddsp/losses_test.py
index c4cb431..0357c1f 100644
--- a/ddsp/losses_test.py
+++ b/ddsp/losses_test.py
@@ -23,10 +23,19 @@ import tensorflow.compat.v2 as tf
 class SpectralLossTest(tf.test.TestCase):
 
   def test_output_shape_is_correct(self):
-    loss_obj = losses.SpectralLoss()
-
-    input_audio = tf.random.uniform((3, 16000), dtype=tf.float32)
-    target_audio = tf.random.uniform((3, 16000), dtype=tf.float32)
+    """Test correct shape with all losses active."""
+    loss_obj = losses.SpectralLoss(
+        mag_weight=1.0,
+        delta_time_weight=1.0,
+        delta_delta_time_weight=1.0,
+        delta_freq_weight=1.0,
+        delta_delta_freq_weight=1.0,
+        logmag_weight=1.0,
+        loudness_weight=1.0,
+    )
+
+    input_audio = tf.random.uniform((3, 8000), dtype=tf.float32)
+    target_audio = tf.random.uniform((3, 8000), dtype=tf.float32)
 
     loss = loss_obj(input_audio, target_audio)
 
diff --git a/ddsp/spectral_ops_test.py b/ddsp/spectral_ops_test.py
index ee40198..fd14940 100644
--- a/ddsp/spectral_ops_test.py
+++ b/ddsp/spectral_ops_test.py
@@ -40,6 +40,25 @@ class STFTTest(tf.test.TestCase):
     self.assertAllClose(np.abs(s_np), np.abs(s_tf), rtol=1e-3, atol=1e-3)
 
 
+class DiffTest(tf.test.TestCase):
+
+  def test_shape_is_correct(self):
+    n_batch = 2
+    n_time = 125
+    n_freq = 100
+    mag = tf.ones([n_batch, n_time, n_freq])
+
+    diff = spectral_ops.diff
+    delta_t = diff(mag, axis=1)
+    self.assertEqual(delta_t.shape[1], mag.shape[1]-1)
+    delta_delta_t = diff(delta_t, axis=1)
+    self.assertEqual(delta_delta_t.shape[1], mag.shape[1]-2)
+    delta_f = diff(mag, axis=2)
+    self.assertEqual(delta_f.shape[2], mag.shape[2]-1)
+    delta_delta_f = diff(delta_f, axis=2)
+    self.assertEqual(delta_delta_f.shape[2], mag.shape[2]-2)
+
+
 class LoudnessTest(tf.test.TestCase):
 
   def test_tf_and_np_are_consistent(self):

pip install -e .[data_preparation,test]
parse_log_pytest
requirements.txt
pytest pylint
apt-get update apt-get install -y gcc apt-get install -y ffmpeg
3.9
requirements/base.txt
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
absl-py==2.2.1
apache-beam==2.63.0
array_record==0.5.1
astroid==3.3.9
astunparse==1.6.3
async-timeout==5.0.1
attrs==25.3.0
audioread==3.0.1
avro-python3==1.10.2
certifi==2025.1.31
cffi==1.17.1
charset-normalizer==3.4.1
click==8.1.8
cloudpickle==2.2.1
contourpy==1.3.0
crcmod==1.7
crepe==0.0.16
cycler==0.12.1
-e git+https://github.com/magenta/ddsp.git@18147a0120f69bf62908c0ae19f3a66735c39b35#egg=ddsp
decorator==5.2.1
dill==0.3.1.1
dm-tree==0.1.8
dnspython==2.7.0
docopt==0.6.2
etils==1.5.2
exceptiongroup==1.2.2
fastavro==1.10.0
fasteners==0.19
flatbuffers==25.2.10
fonttools==4.56.0
fsspec==2025.3.1
future==1.0.0
gast==0.6.0
gin-config==0.5.0
google-pasta==0.2.0
grpcio==1.65.5
h5py==3.13.0
hdfs==2.7.3
hmmlearn==0.3.3
httplib2==0.22.0
idna==3.10
imageio==2.37.0
importlib_metadata==8.6.1
importlib_resources==6.5.2
iniconfig==2.1.0
isort==6.0.1
joblib==1.4.2
jsonpickle==3.4.2
jsonschema==4.23.0
jsonschema-specifications==2024.10.1
keras==3.9.1
kiwisolver==1.4.7
lazy_loader==0.4
libclang==18.1.1
librosa==0.11.0
llvmlite==0.43.0
Markdown==3.7
markdown-it-py==3.0.0
MarkupSafe==3.0.2
matplotlib==3.9.4
mccabe==0.7.0
mdurl==0.1.2
ml_dtypes==0.5.1
msgpack==1.1.0
namex==0.0.8
numba==0.60.0
numpy==2.0.2
objsize==0.7.1
opt_einsum==3.4.0
optree==0.14.1
orjson==3.10.16
packaging==24.2
pillow==11.1.0
platformdirs==4.3.7
pluggy==1.5.0
pooch==1.8.2
promise==2.3
proto-plus==1.26.1
protobuf==3.20.3
psutil==7.0.0
pyarrow==16.1.0
pyarrow-hotfix==0.6
pycparser==2.22
pydot==1.4.2
pydub==0.25.1
Pygments==2.19.1
pylint==3.3.6
pymongo==4.11.3
pyparsing==3.2.3
pytest==8.3.5
python-dateutil==2.9.0.post0
pytz==2025.2
PyYAML==6.0.2
redis==5.2.1
referencing==0.36.2
regex==2024.11.6
requests==2.32.3
resampy==0.4.3
rich==14.0.0
rpds-py==0.24.0
scikit-learn==1.6.1
scipy==1.13.1
six==1.17.0
sortedcontainers==2.4.0
soundfile==0.13.1
soxr==0.5.0.post1
tensorboard==2.19.0
tensorboard-data-server==0.7.2
tensorflow==2.19.0
tensorflow-io-gcs-filesystem==0.37.1
tensorflow-metadata==1.16.1
termcolor==2.5.0
tfds-nightly==4.9.3.dev202311230044
threadpoolctl==3.6.0
toml==0.10.2
tomli==2.2.1
tomlkit==0.13.2
tqdm==4.67.1
typing_extensions==4.13.0
urllib3==2.3.0
Werkzeug==3.1.3
wrapt==1.17.2
zipp==3.21.0
zstandard==0.23.0

name: ddsp
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=3.0.16=h5eee18b_0
  - pip=25.0=py39h06a4308_0
  - python=3.9.21=he870216_1
  - readline=8.2=h5eee18b_0
  - setuptools=75.8.0=py39h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - tzdata=2025a=h04d1e81_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - absl-py==2.2.1
      - apache-beam==2.63.0
      - array-record==0.5.1
      - astroid==3.3.9
      - astunparse==1.6.3
      - async-timeout==5.0.1
      - attrs==25.3.0
      - audioread==3.0.1
      - avro-python3==1.10.2
      - certifi==2025.1.31
      - cffi==1.17.1
      - charset-normalizer==3.4.1
      - click==8.1.8
      - cloudpickle==2.2.1
      - contourpy==1.3.0
      - crcmod==1.7
      - crepe==0.0.16
      - cycler==0.12.1
      - decorator==5.2.1
      - dill==0.3.1.1
      - dm-tree==0.1.8
      - dnspython==2.7.0
      - docopt==0.6.2
      - etils==1.5.2
      - exceptiongroup==1.2.2
      - fastavro==1.10.0
      - fasteners==0.19
      - flatbuffers==25.2.10
      - fonttools==4.56.0
      - fsspec==2025.3.1
      - future==1.0.0
      - gast==0.6.0
      - gin-config==0.5.0
      - google-pasta==0.2.0
      - grpcio==1.65.5
      - h5py==3.13.0
      - hdfs==2.7.3
      - hmmlearn==0.3.3
      - httplib2==0.22.0
      - idna==3.10
      - imageio==2.37.0
      - importlib-metadata==8.6.1
      - importlib-resources==6.5.2
      - iniconfig==2.1.0
      - isort==6.0.1
      - joblib==1.4.2
      - jsonpickle==3.4.2
      - jsonschema==4.23.0
      - jsonschema-specifications==2024.10.1
      - keras==3.9.1
      - kiwisolver==1.4.7
      - lazy-loader==0.4
      - libclang==18.1.1
      - librosa==0.11.0
      - llvmlite==0.43.0
      - markdown==3.7
      - markdown-it-py==3.0.0
      - markupsafe==3.0.2
      - matplotlib==3.9.4
      - mccabe==0.7.0
      - mdurl==0.1.2
      - ml-dtypes==0.5.1
      - msgpack==1.1.0
      - namex==0.0.8
      - numba==0.60.0
      - numpy==2.0.2
      - objsize==0.7.1
      - opt-einsum==3.4.0
      - optree==0.14.1
      - orjson==3.10.16
      - packaging==24.2
      - pillow==11.1.0
      - platformdirs==4.3.7
      - pluggy==1.5.0
      - pooch==1.8.2
      - promise==2.3
      - proto-plus==1.26.1
      - protobuf==3.20.3
      - psutil==7.0.0
      - pyarrow==16.1.0
      - pyarrow-hotfix==0.6
      - pycparser==2.22
      - pydot==1.4.2
      - pydub==0.25.1
      - pygments==2.19.1
      - pylint==3.3.6
      - pymongo==4.11.3
      - pyparsing==3.2.3
      - pytest==8.3.5
      - python-dateutil==2.9.0.post0
      - pytz==2025.2
      - pyyaml==6.0.2
      - redis==5.2.1
      - referencing==0.36.2
      - regex==2024.11.6
      - requests==2.32.3
      - resampy==0.4.3
      - rich==14.0.0
      - rpds-py==0.24.0
      - scikit-learn==1.6.1
      - scipy==1.13.1
      - six==1.17.0
      - sortedcontainers==2.4.0
      - soundfile==0.13.1
      - soxr==0.5.0.post1
      - tensorboard==2.19.0
      - tensorboard-data-server==0.7.2
      - tensorflow==2.19.0
      - tensorflow-io-gcs-filesystem==0.37.1
      - tensorflow-metadata==1.16.1
      - termcolor==2.5.0
      - tfds-nightly==4.9.3.dev202311230044
      - threadpoolctl==3.6.0
      - toml==0.10.2
      - tomli==2.2.1
      - tomlkit==0.13.2
      - tqdm==4.67.1
      - typing-extensions==4.13.0
      - urllib3==2.3.0
      - werkzeug==3.1.3
      - wrapt==1.17.2
      - zipp==3.21.0
      - zstandard==0.23.0
prefix: /opt/conda/envs/ddsp

ddsp/losses_test.py::SpectralLossTest::test_output_shape_is_correct ddsp/spectral_ops_test.py::DiffTest::test_shape_is_correct
ddsp/losses_test.py::PretrainedCREPEEmbeddingLossTest::test_output_shape_is_correct ddsp/spectral_ops_test.py::STFTTest::test_tf_and_np_are_consistent ddsp/spectral_ops_test.py::LoudnessTest::test_tf_and_np_are_consistent
ddsp/losses.py ddsp/version.py ddsp/spectral_ops.py
ddsp/losses.py ddsp/version.py ddsp/spectral_ops.py