diff --git a/codes/models/diffusion/diffusion_dvae.py b/codes/models/diffusion/diffusion_dvae.py
index fa1f9b50..ea663826 100644
--- a/codes/models/diffusion/diffusion_dvae.py
+++ b/codes/models/diffusion/diffusion_dvae.py
@@ -251,11 +251,7 @@ class DiffusionDVAE(nn.Module):
         )
 
     def get_debug_values(self, step, __):
-        if self.record_codes:
-            # Report annealing schedule
-            return {'histogram_codes': self.codes}
-        else:
-            return {}
+        return {'histogram_codes': self.codes}
 
     @torch.no_grad()
     @eval_decorator
@@ -379,10 +375,8 @@ class DiffusionDVAE(nn.Module):
 
 # Test for ~4 second audio clip at 22050Hz
 if __name__ == '__main__':
-    spec = torch.randn(4, 80, 416)
-    cond = torch.randn(4, 5, 80, 200)
-    num_cond = torch.tensor([2,4,5,3], dtype=torch.long)
+    spec = torch.randn(4, 80, 161)
     ts = torch.LongTensor([432, 234, 100, 555])
     model = DiffusionDVAE(model_channels=128, num_res_blocks=1, in_channels=80, out_channels=160, spectrogram_conditioning_levels=[1,2],
     channel_mult=(1,2,4), attention_resolutions=[4], num_heads=4, kernel_size=3, scale_steps=2, conditioning_inputs_provided=False)
-    print(model(torch.randn_like(spec), ts, spec, cond, num_cond)[0].shape)
+    print(model(torch.randn_like(spec), ts, spec)[0].shape)
diff --git a/codes/trainer/injectors/base_injectors.py b/codes/trainer/injectors/base_injectors.py
index 90a7ba69..02d20b7d 100644
--- a/codes/trainer/injectors/base_injectors.py
+++ b/codes/trainer/injectors/base_injectors.py
@@ -565,5 +565,5 @@ class RandomAudioCropInjector(Injector):
 
 
 if __name__ == '__main__':
-    inj = DecomposeDimensionInjector({'dim':2, 'in': 'x', 'out': 'y'}, None)
-    print(inj({'x':torch.randn(10,3,64,64)})['y'].shape)
\ No newline at end of file
+    inj = MelSpectrogramInjector({'in': 'x', 'out': 'y'}, None)
+    print(inj({'x':torch.rand(10,1,40800)})['y'].shape)
\ No newline at end of file