-
Notifications
You must be signed in to change notification settings - Fork 967
/
Copy.lua
42 lines (33 loc) · 1.15 KB
/
Copy.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
local Copy, parent = torch.class('nn.Copy', 'nn.Module')
function Copy:__init(intype, outtype, forceCopy, dontCast)
intype = intype or torch.Tensor.__typename
outtype = outtype or torch.Tensor.__typename
self.dontCast = dontCast
parent.__init(self)
self.gradInput = torch.getmetatable(intype).new()
self.output = torch.getmetatable(outtype).new()
if (not forceCopy) and intype == outtype then
self.updateOutput = function(self, input)
self.output:set(input)
return input
end
self.updateGradInput = function(self, input, gradOutput)
self.gradInput:set(gradOutput)
return gradOutput
end
end
end
function Copy:updateOutput(input)
self.output:resize(input:size()):copy(input)
return self.output
end
function Copy:updateGradInput(input, gradOutput)
self.gradInput:resize(gradOutput:size()):copy(gradOutput)
return self.gradInput
end
function Copy:type(type, tensorCache)
if type and self.dontCast then
return self
end
return parent.type(self, type, tensorCache)
end