Update model.py
Browse files
model.py
CHANGED
|
@@ -121,6 +121,7 @@ class VAE(tf.keras.Model):
|
|
| 121 |
self.encoder = Encoder_Z(dim_z=self.dim_z).build()
|
| 122 |
self.decoder = Decoder_X(dim_z=self.dim_z).build()
|
| 123 |
self.cost_func = VAECost(self)
|
|
|
|
| 124 |
|
| 125 |
@tf.function()
|
| 126 |
def train_step(self, data):
|
|
|
|
| 121 |
self.encoder = Encoder_Z(dim_z=self.dim_z).build()
|
| 122 |
self.decoder = Decoder_X(dim_z=self.dim_z).build()
|
| 123 |
self.cost_func = VAECost(self)
|
| 124 |
+
self.load_weights("weights/")
|
| 125 |
|
| 126 |
@tf.function()
|
| 127 |
def train_step(self, data):
|