add register_to_config, pass training tests

Author: natolambert

src/diffusers/models/unet_rl.py

@@ -5,7 +5,7 @@
 
 from diffusers.models.resnet import Downsample1D, ResidualTemporalBlock, Upsample1D
 
-from ..configuration_utils import ConfigMixin
+from ..configuration_utils import ConfigMixin, register_to_config
 from ..modeling_utils import ModelMixin
 from .embeddings import get_timestep_embedding
 
@@ -57,6 +57,7 @@ def forward(self, x):
 
 
 class TemporalUNet(ModelMixin, ConfigMixin):  # (nn.Module):
+    @register_to_config
     def __init__(
         self,
         training_horizon=128,

tests/test_modeling_utils.py

@@ -629,6 +629,12 @@ def input_shape(self):
     def output_shape(self):
         return (4, 16, 14)
 
+    def test_ema_training(self):
+        pass
+
+    def test_training(self):
+        pass
+
     def prepare_init_args_and_inputs_for_common(self):
         init_dict = {
             "training_horizon": 128,