diff --git a/hpc/test/testmodel.py b/hpc/test/testmodel.py new file mode 100644 index 0000000..ca44d70 --- /dev/null +++ b/hpc/test/testmodel.py @@ -0,0 +1,34 @@ +import umbridge +import time +import os + +class TestModel(umbridge.Model): + + def __init__(self): + super().__init__("forward") + + def get_input_sizes(self, config): + return [1] + + def get_output_sizes(self, config): + return [1] + + def __call__(self, parameters, config): + # Sleep for number of milliseconds defined in env var + time.sleep(int(os.getenv("TEST_DELAY", 0)) / 1000) + + posterior = 2*parameters[0][0] + return [[posterior]] + + def supports_evaluate(self): + return True + + def gradient(self,out_wrt, in_wrt, parameters, sens, config): + # Sleep for number of milliseconds defined in env var + time.sleep(int(os.getenv("TEST_DELAY", 0)) / 1000) + return [2*sens[0]] + + def supports_gradient(self): + return True + +umbridge.serve_models([TestModel()], int(os.getenv("PORT", 4242)))