MLA_Learn_Periodic_Functions/code/RNN/Model.py

import tensorflow as tf

class MyModel(tf.keras.Model):

	def __init__(self, HIDDEN):
		super(MyModel, self).__init__()
		self.lstm1 = tf.keras.layers.LSTM(HIDDEN, return_sequences=True)
		self.lstm2 = tf.keras.layers.LSTM(HIDDEN, return_sequences=True)
		#self.lstm3 = tf.keras.layers.LSTM(HIDDEN, return_sequences=True)
		#self.lstm4 = tf.keras.layers.LSTM(HIDDEN, return_sequences=True)
		self.lstmlast = tf.keras.layers.LSTM(1, return_sequences=True)


	def call(self, inputs):
		x = self.lstm1(inputs)
		x = self.lstm2(x)
		#x = self.lstm3(x)
		#x = self.lstm4(x)
		x = self.lstmlast(x)
		return x