kerasで損失関数を自作

やりたいこと

ある9×120行列X=x_j(j=1,…,9)に対して、損失関数(目的関数)
-1/9Σ[j=1..9][-1/0.8exp(-0.8I(X))]
を最小化するようなI(X)の関数形をニューラルネットワークで導出しようとしています。

言うなれば、Xをニューラルネットワークに入力した時、出力されるI(X)は活性化関数fと重みωを用いて
I(X)=f(ωX)
と表されるため(実際は中間層があるのでもっと複雑ですが)、その重みωを求めたいといった感じです。

該当のソースコード

import numpy as np
import pandas as pd
import csv

from keras.models import Sequential
from keras.layers import Dense, Activation
from pandas import DataFrame

import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.losses import Loss

import os

#CSVファイルの読み込み
csvdata = pd.read_csv("行列.csv", header=None)
x_train = DataFrame(csvdata)

#CSVから読んだデータをnumpyの行列に入れる
myarray = np.array
myarray = csvdata.values
x_train = np.asarray(x_train).astype(np.float32)

#乱数シード固定
import random
def set_seed(seed=200):
    tf.random.set_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)

#ニューラルネットワーク作成(64ノードの中間層1つ)
set_seed(0)
model = Sequential()
model.add(Dense(64, input_dim=120))
model.add(Activation('relu'))
model.add(Dense(1, input_dim=64))
model.add(Activation('relu'))

#損失関数作成
def custom_loss(y_pred):
    U = -np.exp(-0.8*y_pred)/0.8
    return K.mean( U )

#自作した損失関数クラスを実装
class CustomLoss(Loss):
  def __init__(self, name="custom_loss", **kwargs):
    super(CustomLoss, self).__init__(name=name, **kwargs)

def call(self, y_pred):
    y_pred = tf.convert_to_tensor(y_pred)
    return custom_loss(y_pred)
    #念のためTensor化

#最適化
model.compile(optimizer='SGD', loss=custom_loss, metrics=['accuracy'])

#重みを出力
print(model.layers[0].get_weights()[0])