From: Francois Fleuret <francois@fleuret.org>
Date: Tue, 6 Dec 2016 16:33:37 +0000 (+0100)
Subject: Decorate as high as possible in the class hierarchy.
X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=commitdiff_plain;h=7c5787015c28b2c551bbd1016512fcb8baf5d3e8;p=profiler-torch.git

Decorate as high as possible in the class hierarchy.
---

diff --git a/profiler.lua b/profiler.lua
index f57c9d8..77461b3 100644
--- a/profiler.lua
+++ b/profiler.lua
@@ -64,15 +64,20 @@ function profiler.decorate(model, functionsToDecorate)
       -- save models anymore.
 
       if rawget(model, name) then
-         error('We decorate the class, not the objects, and there is a `'
+         error('We decorate the classes, not the objects, and there is a `'
                   .. name
                   .. '\' function in '
                   .. tostring(model))
       end
 
-      local toDecorate = getmetatable(model)
+      local toDecorate = model
 
-      if toDecorate[name] and not toDecorate[nameOrig] then
+      while not rawget(toDecorate, name) do
+         toDecorate = getmetatable(toDecorate)
+      end
+
+      if not toDecorate[nameOrig] then
+         print('Decorating ' .. toDecorate.__typename .. '.' .. name)
          toDecorate[nameOrig] = toDecorate[name]
          toDecorate[name] = function(self, ...)
             local startTime = sys.clock()